1package dgram
2
3import (
4 "bytes"
5 "crypto/sha256"
6 "encoding/base64"
7 "encoding/json"
8 "os"
9 "path/filepath"
10 "testing"
11)
12
13func privateDir(t *testing.T) string {
14 t.Helper()
15 return filepath.Join(t.TempDir(), "vev")
16}
17
18func TestKeyFingerprint(t *testing.T) {
19 key := bytes.Repeat([]byte{0xA5}, 32)
20 sum := sha256.Sum256(key)
21 want := base64.StdEncoding.EncodeToString(sum[:])
22 if got := KeyFingerprint(key); got != want {
23 t.Fatalf("KeyFingerprint()=%q, want %q", got, want)
24 }
25 other := bytes.Repeat([]byte{0xA6}, 32)
26 if KeyFingerprint(key) == KeyFingerprint(other) {
27 t.Fatal("distinct keys produced identical fingerprints")
28 }
29}
30
31func TestProxyRegistryFileHoldsNoKeyMaterial(t *testing.T) {
32 r := NewProxyRegistry(privateDir(t))
33 key := bytes.Repeat([]byte{0xA5}, 32)
34 fingerprint := KeyFingerprint(key)
35 rec := ProxyRecord{Session: "work", PID: os.Getpid(), Port: 61000, KeyFingerprint: fingerprint}
36 if err := r.Publish(rec); err != nil {
37 t.Fatal(err)
38 }
39
40 data, err := os.ReadFile(r.path("work"))
41 if err != nil {
42 t.Fatal(err)
43 }
44 if bytes.Contains(data, key) || bytes.Contains(data, []byte(base64.StdEncoding.EncodeToString(key))) {
45 t.Fatalf("registry file contains raw key material: %s", data)
46 }
47 var fields map[string]json.RawMessage
48 if err := json.Unmarshal(data, &fields); err != nil {
49 t.Fatal(err)
50 }
51 if _, ok := fields["key"]; ok {
52 t.Fatalf("registry file contains raw key field: %s", data)
53 }
54 var got string
55 if err := json.Unmarshal(fields["key_fingerprint"], &got); err != nil {
56 t.Fatal(err)
57 }
58 if got != fingerprint {
59 t.Fatalf("key_fingerprint=%q, want %q", got, fingerprint)
60 }
61}
62
63func TestProxyRegistryReuseStaleCleanupAndSupersede(t *testing.T) {
64 t.Run("reuses live matching record", func(t *testing.T) {
65 r := NewProxyRegistry(privateDir(t))
66 want := ProxyRecord{Session: "work", PID: os.Getpid(), Port: 61000, KeyFingerprint: "abc"}
67 if err := r.Publish(want); err != nil {
68 t.Fatal(err)
69 }
70 got, ok := r.Lookup("work")
71 if !ok {
72 t.Fatal("Lookup ok=false, want true")
73 }
74 if got.Port != want.Port || got.KeyFingerprint != want.KeyFingerprint || got.PID != want.PID {
75 t.Fatalf("Lookup()=%+v, want %+v", got, want)
76 }
77 })
78 t.Run("removes stale record", func(t *testing.T) {
79 r := NewProxyRegistry(privateDir(t))
80 if err := r.Publish(ProxyRecord{Session: "work", PID: os.Getpid(), Port: 61000, KeyFingerprint: "abc"}); err != nil {
81 t.Fatal(err)
82 }
83 if err := os.WriteFile(r.path("work"), []byte(`{"session":"work","pid":-1,"port":61000,"key_fingerprint":"abc"}`), 0o600); err != nil {
84 t.Fatal(err)
85 }
86 if _, ok := r.Lookup("work"); ok {
87 t.Fatal("Lookup ok=true for stale pid")
88 }
89 if _, err := os.Stat(r.path("work")); !os.IsNotExist(err) {
90 t.Fatalf("stale record stat err=%v, want not exist", err)
91 }
92 })
93 t.Run("supersedes without pid-based kill and cleanup is owned", func(t *testing.T) {
94 r := NewProxyRegistry(privateDir(t))
95 old := ProxyRecord{Session: "work", PID: os.Getpid(), Port: 61000, KeyFingerprint: "old"}
96 if err := r.Publish(old); err != nil {
97 t.Fatal(err)
98 }
99 newRec := ProxyRecord{Session: "work", PID: os.Getpid(), Port: 61001, KeyFingerprint: "new"}
100 if err := r.Publish(newRec); err != nil {
101 t.Fatal(err)
102 }
103 if err := r.RemoveOwned(old); err != nil {
104 t.Fatal(err)
105 }
106 got, ok := r.Lookup("work")
107 if !ok {
108 t.Fatal("new live record removed by stale owner")
109 }
110 if got.Port != newRec.Port || got.KeyFingerprint != newRec.KeyFingerprint {
111 t.Fatalf("Lookup()=%+v, want %+v", got, newRec)
112 }
113 if err := r.RemoveOwned(newRec); err != nil {
114 t.Fatal(err)
115 }
116 if _, ok := r.Lookup("work"); ok {
117 t.Fatal("owned remove left registry record")
118 }
119 })
120}
121
122func TestProxyRegistrySessionPathIsCollisionFree(t *testing.T) {
123 r := NewProxyRegistry(privateDir(t))
124 tests := []struct {
125 a string
126 b string
127 }{
128 {a: "", b: "default"},
129 {a: "a/b", b: "a_b"},
130 {a: `a\b`, b: "a_b"},
131 }
132 for _, tt := range tests {
133 if r.path(tt.a) == r.path(tt.b) {
134 t.Fatalf("path(%q) collides with path(%q): %s", tt.a, tt.b, r.path(tt.a))
135 }
136 first := ProxyRecord{Session: tt.a, PID: os.Getpid(), Port: 61000, KeyFingerprint: "first"}
137 second := ProxyRecord{Session: tt.b, PID: os.Getpid(), Port: 61001, KeyFingerprint: "second"}
138 if err := r.Publish(first); err != nil {
139 t.Fatal(err)
140 }
141 if err := r.Publish(second); err != nil {
142 t.Fatal(err)
143 }
144 gotFirst, ok := r.Lookup(tt.a)
145 if !ok || gotFirst.KeyFingerprint != first.KeyFingerprint {
146 t.Fatalf("Lookup(%q)=%+v ok=%v, want first record", tt.a, gotFirst, ok)
147 }
148 gotSecond, ok := r.Lookup(tt.b)
149 if !ok || gotSecond.KeyFingerprint != second.KeyFingerprint {
150 t.Fatalf("Lookup(%q)=%+v ok=%v, want second record", tt.b, gotSecond, ok)
151 }
152 }
153}