DriverVault/API Server/internal/plugins/builtin/ankersolix/payloadcrypto_test.go

package ankersolix

import (
	"context"
	"crypto/ecdh"
	"crypto/hmac"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
)

// encryptedCloud is the server half of the scheme in payloadcrypto.go, written
// from the reference implementation's description rather than from our client,
// so a test against it checks the two halves agree on every step: the presetKey
// wrapping of both public keys, the handshake signature, the per-request
// signature under the session key, the sealed body and the sealed reply.
type encryptedCloud struct {
	t *testing.T

	mu        sync.Mutex
	secrets   map[string][]byte // x-key-ident -> shared secret
	logins    int
	exchanges int
	seen      []map[string]any // decrypted request bodies, in order
	rejectTok string           // a token the cloud answers 401 to
}

func newEncryptedCloud(t *testing.T) *encryptedCloud {
	return &encryptedCloud{t: t, secrets: map[string][]byte{}}
}

func (c *encryptedCloud) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	c.mu.Lock()
	defer c.mu.Unlock()
	body, _ := io.ReadAll(r.Body)

	switch {
	case strings.HasSuffix(r.URL.Path, epLogin):
		c.logins++
		tok := "tok-" + hex.EncodeToString([]byte{byte(c.logins)})
		_, _ = w.Write([]byte(`{"code":0,"data":{"auth_token":"` + tok + `","user_id":"u1","nick_name":"n","token_expires_at":9999999999}}`))

	case strings.HasSuffix(r.URL.Path, epKeyExchange):
		c.exchanges++
		if got := r.Header.Get("content-type"); got != "application/json" {
			c.t.Errorf("handshake content-type = %q, want application/json", got)
		}
		if r.Header.Get("x-auth-ts") == "" || r.Header.Get("x-auth-token") == "" {
			c.t.Error("handshake must carry x-auth-ts and the auth token")
		}
		var req struct {
			ClientPublicKey string `json:"client_public_key"`
		}
		if err := json.Unmarshal(body, &req); err != nil || req.ClientPublicKey == "" {
			c.t.Errorf("handshake body %q is not {client_public_key}", body)
		}
		c.checkSignature(r, presetKeyHex, req.ClientPublicKey)

		preset, _ := hex.DecodeString(presetKeyHex)
		pointHex, err := aesCBCOpen(preset, req.ClientPublicKey)
		if err != nil {
			c.t.Fatalf("unwrap client key: %v", err)
		}
		point, _ := hex.DecodeString(string(pointHex))
		if len(point) != 65 || point[0] != 4 {
			c.t.Fatalf("client key is %d bytes starting %x, want a 65-byte uncompressed point", len(point), point[:1])
		}
		clientPub, err := ecdh.P256().NewPublicKey(point)
		if err != nil {
			c.t.Fatalf("client key: %v", err)
		}
		serverPriv, _ := ecdh.P256().GenerateKey(rand.Reader)
		secret, _ := serverPriv.ECDH(clientPub)
		c.secrets[r.Header.Get("x-key-ident")] = secret
		wrapped, _ := aesCBCSeal(preset, []byte(hex.EncodeToString(serverPriv.PublicKey().Bytes())))
		_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "data": map[string]string{"server_public_key": wrapped}})

	default:
		if r.Header.Get("x-auth-token") == c.rejectTok {
			w.WriteHeader(http.StatusUnauthorized)
			return
		}
		if got := r.Header.Get("content-type"); got != "text/plain" {
			c.t.Errorf("%s content-type = %q, want text/plain", r.URL.Path, got)
		}
		if got := r.Header.Get("x-encryption-info"); got != "algo_ecdh" {
			c.t.Errorf("x-encryption-info = %q", got)
		}
		secret := c.secrets[r.Header.Get("x-key-ident")]
		if secret == nil {
			c.t.Errorf("%s arrived under key ident %q the cloud never issued", r.URL.Path, r.Header.Get("x-key-ident"))
			w.WriteHeader(http.StatusBadRequest)
			return
		}
		c.checkSignature(r, hex.EncodeToString(secret[:16]), string(body))
		plain, err := aesCBCOpen(secret[:16], string(body))
		if err != nil {
			c.t.Fatalf("open %s body: %v", r.URL.Path, err)
		}
		var payload map[string]any
		if err := json.Unmarshal(plain, &payload); err != nil {
			c.t.Fatalf("%s body decrypted to %q, not JSON", r.URL.Path, plain)
		}
		c.seen = append(c.seen, payload)

		reply, _ := json.Marshal(map[string]any{"echo": payload, "path": r.URL.Path, "query": r.URL.RawQuery})
		sealed, _ := aesCBCSeal(secret[:16], reply)
		_ = json.NewEncoder(w).Encode(map[string]any{"code": 0, "msg": "success!", "data": sealed, "signature": "sig"})
	}
}

// checkSignature recomputes x-signature the way the reference documents it:
// HMAC-SHA256 keyed on the ascii hex of the key, over ts+once+wire.
func (c *encryptedCloud) checkSignature(r *http.Request, keyHex, wire string) {
	ts, once := r.Header.Get("x-request-ts"), r.Header.Get("x-request-once")
	if len(once) != 32 || len(r.Header.Get("x-key-ident")) != 32 {
		c.t.Errorf("x-request-once %q / x-key-ident %q should be 32 hex chars", once, r.Header.Get("x-key-ident"))
	}
	mac := hmac.New(sha256.New, []byte(keyHex))
	mac.Write([]byte(ts + "+" + once + "+" + wire))
	if want := hex.EncodeToString(mac.Sum(nil)); r.Header.Get("x-signature") != want {
		c.t.Errorf("%s x-signature = %q, want %q", r.URL.Path, r.Header.Get("x-signature"), want)
	}
}

func encryptedPlugin(t *testing.T, srv *httptest.Server) *Plugin {
	t.Helper()
	p := &Plugin{}
	if err := p.Init(context.Background(), map[string]string{"email": "e@example.com", "password": "pw", "payloadEncryption": "on"}); err != nil {
		t.Fatal(err)
	}
	p.apiBase = srv.URL
	p.sess = sessionFor(srv.URL, "e@example.com", "pw")
	return p
}

// TestPayloadEncryptionRoundTrip drives a POST and a GET through the sealed
// path and checks the plugin sees exactly the envelope the plain mode would.
func TestPayloadEncryptionRoundTrip(t *testing.T) {
	cloud := newEncryptedCloud(t)
	srv := httptest.NewServer(cloud)
	defer srv.Close()
	p := encryptedPlugin(t, srv)

	body, err := p.apiRequest(context.Background(), epStationInfo, map[string]any{"sn": "ABC", "featuretype": 1})
	if err != nil {
		t.Fatal(err)
	}
	var env struct {
		Code int `json:"code"`
		Data struct {
			Echo map[string]any `json:"echo"`
			Path string         `json:"path"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &env); err != nil {
		t.Fatalf("unsealed body %s: %v", body, err)
	}
	if env.Code != 0 || env.Data.Echo["sn"] != "ABC" || env.Data.Echo["featuretype"] != float64(1) {
		t.Fatalf("unsealed envelope = %s", body)
	}
	if !strings.HasSuffix(env.Data.Path, epStationInfo) {
		t.Errorf("request went to %q", env.Data.Path)
	}

	body, err = p.apiGet(context.Background(), epMessages, map[string]string{"last_time": "7"})
	if err != nil {
		t.Fatal(err)
	}
	var got struct {
		Data struct {
			Echo  map[string]any `json:"echo"`
			Query string         `json:"query"`
		} `json:"data"`
	}
	_ = json.Unmarshal(body, &got)
	if got.Data.Query != "last_time=7" {
		t.Errorf("GET query = %q, want last_time=7", got.Data.Query)
	}
	if len(got.Data.Echo) != 0 {
		t.Errorf("a GET seals an empty object, got %v", got.Data.Echo)
	}

	if cloud.logins != 1 || cloud.exchanges != 1 {
		t.Errorf("logins = %d, exchanges = %d; want one of each for two requests", cloud.logins, cloud.exchanges)
	}
}

// TestPayloadEncryptionHandshakeFollowsLogin: a rejected token costs a login and
// a new handshake, and the retry goes out under both.
func TestPayloadEncryptionHandshakeFollowsLogin(t *testing.T) {
	cloud := newEncryptedCloud(t)
	srv := httptest.NewServer(cloud)
	defer srv.Close()
	p := encryptedPlugin(t, srv)

	if _, err := p.apiRequest(context.Background(), epSiteList, map[string]any{}); err != nil {
		t.Fatal(err)
	}
	cloud.mu.Lock()
	cloud.rejectTok = "tok-01" // the first login's token
	cloud.mu.Unlock()

	if _, err := p.apiRequest(context.Background(), epSiteList, map[string]any{}); err != nil {
		t.Fatal(err)
	}
	if cloud.logins != 2 || cloud.exchanges != 2 {
		t.Errorf("logins = %d, exchanges = %d; want 2 and 2 after a rejected token", cloud.logins, cloud.exchanges)
	}
	if p.sess.enc == nil || !p.sess.enc.ready() {
		t.Error("the session should hold the new handshake")
	}
}

// TestPayloadEncryptionOffStaysPlain guards the default: no switch, no
// handshake, JSON on the wire as before.
func TestPayloadEncryptionOffStaysPlain(t *testing.T) {
	var sawExchange bool
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch {
		case strings.HasSuffix(r.URL.Path, epLogin):
			_, _ = w.Write([]byte(`{"code":0,"data":{"auth_token":"t","user_id":"u1","token_expires_at":9999999999}}`))
		case strings.HasSuffix(r.URL.Path, epKeyExchange):
			sawExchange = true
		default:
			if r.Header.Get("content-type") != "application/json" || r.Header.Get("x-encryption-info") != "" {
				t.Errorf("plain request carried %q / %q", r.Header.Get("content-type"), r.Header.Get("x-encryption-info"))
			}
			_, _ = w.Write([]byte(`{"code":0,"data":{}}`))
		}
	}))
	defer srv.Close()

	p := &Plugin{}
	_ = p.Init(context.Background(), map[string]string{"email": "plain@example.com", "password": "pw"})
	p.apiBase = srv.URL
	p.sess = sessionFor(srv.URL, "plain@example.com", "pw")
	if p.encrypt {
		t.Fatal("payload encryption must be off by default")
	}
	if _, err := p.apiRequest(context.Background(), epSiteList, map[string]any{}); err != nil {
		t.Fatal(err)
	}
	if sawExchange {
		t.Error("no handshake should run with encryption off")
	}
}

// TestUnsealEnvelopeLeavesUnsignedAlone: an error envelope, or a plain reply,
// is passed through untouched; a signed reply that fails to open is an error
// rather than garbage.
func TestUnsealEnvelopeLeavesUnsignedAlone(t *testing.T) {
	enc, _ := newPayloadCipher()
	enc.secret = make([]byte, 32)

	for _, in := range []string{`{"code":26161,"msg":"Failed to request."}`, `{"code":0,"data":{"x":1}}`, `not json`} {
		out, err := enc.unsealEnvelope([]byte(in))
		if err != nil || string(out) != in {
			t.Errorf("unsealEnvelope(%q) = %q, %v; want it untouched", in, out, err)
		}
	}
	if _, err := enc.unsealEnvelope([]byte(`{"code":0,"data":"AAAA","signature":"s"}`)); err == nil {
		t.Error("a signed reply that does not open should be an error")
	}
}

func TestPayloadEncryptionSwitch(t *testing.T) {
	for v, want := range map[string]bool{"": false, "off": false, "on": true, "ON": true, "true": true, "1": true, "no": false} {
		if got := payloadEncryptionOn(v); got != want {
			t.Errorf("payloadEncryptionOn(%q) = %v", v, got)
		}
	}
	fields := map[string]bool{}
	for _, f := range (&Plugin{}).Descriptor().ConfigFields {
		fields[f.Key] = f.Default == "off"
	}
	if off, ok := fields["payloadEncryption"]; !ok || !off {
		t.Error("payloadEncryption must be advertised with default off")
	}
}