1package pico
  2
  3import (
  4	"bytes"
  5	"crypto/ed25519"
  6	"fmt"
  7	"io"
  8	"log/slog"
  9	"net"
 10	"strings"
 11	"testing"
 12
 13	"github.com/picosh/pico/pkg/db"
 14	"github.com/picosh/pico/pkg/shared"
 15	"github.com/picosh/pico/pkg/shared/router"
 16	"golang.org/x/crypto/ssh"
 17)
 18
 19type fakeSession struct{ bytes.Buffer }
 20
 21func (s *fakeSession) Exit(int) error        { return nil }
 22func (s *fakeSession) Close() error          { return nil }
 23func (s *fakeSession) Stderr() io.ReadWriter { return &s.Buffer }
 24
 25type dnsCheckDB struct {
 26	db.DB
 27	users    []*db.User
 28	projects map[string]bool
 29	keys     []*db.PublicKey
 30}
 31
 32func (d *dnsCheckDB) FindKeysByUser(user *db.User) ([]*db.PublicKey, error) {
 33	keys := []*db.PublicKey{}
 34	for _, k := range d.keys {
 35		if k.UserID == user.ID {
 36			keys = append(keys, k)
 37		}
 38	}
 39	return keys, nil
 40}
 41
 42func (d *dnsCheckDB) FindUserByName(name string) (*db.User, error) {
 43	for _, u := range d.users {
 44		if u.Name == name {
 45			return u, nil
 46		}
 47	}
 48	return nil, fmt.Errorf("user not found")
 49}
 50
 51func (d *dnsCheckDB) FindProjectByName(userID, name string) (*db.Project, error) {
 52	if d.projects[userID+"/"+name] {
 53		return &db.Project{UserID: userID, Name: name}, nil
 54	}
 55	return nil, fmt.Errorf("project not found")
 56}
 57
 58func TestDNSCheck(t *testing.T) {
 59	me := &db.User{ID: "1", Name: "antonio"}
 60	other := &db.User{ID: "2", Name: "erock"}
 61	edKey, _, err := ed25519.GenerateKey(nil)
 62	if err != nil {
 63		t.Fatal(err)
 64	}
 65	pk, err := ssh.NewPublicKey(edKey)
 66	if err != nil {
 67		t.Fatal(err)
 68	}
 69	fingerprint := ssh.FingerprintSHA256(pk)
 70	dbpool := &dnsCheckDB{
 71		users:    []*db.User{me, other},
 72		projects: map[string]bool{"1/blog": true},
 73		keys:     []*db.PublicKey{{UserID: me.ID, Name: "laptop", Key: shared.KeyForKeyText(pk)}},
 74	}
 75
 76	tests := []struct {
 77		name    string
 78		domain  string
 79		records map[string][]string
 80		want    []string
 81	}{
 82		{"no records", "", map[string][]string{}, []string{
 83			"_pgs.example.com: no TXT record",
 84			"_prose.example.com: no TXT record",
 85			"_tuns.example.com: no TXT record",
 86			"_sish.example.com: no TXT record",
 87		}},
 88		{"valid", "", map[string][]string{
 89			"pgs":   {"antonio-blog"},
 90			"prose": {"erock"},
 91			"sish":  {"SHA256:someone-else", fingerprint},
 92		}, []string{
 93			`_pgs.example.com: "antonio-blog" is valid, serving pgs project "blog" of antonio`,
 94			`_prose.example.com: "erock" is valid, serving prose blog of erock (not your account)`,
 95			fmt.Sprintf(`_sish.example.com: %q is valid, tunnels from your key "laptop" can use this domain`, fingerprint),
 96		}},
 97		{"invalid", "", map[string][]string{
 98			"pgs":   {"antonio-blgo"},
 99			"prose": {"nobody"},
100			"sish":  {"SHA256:someone-else"},
101		}, []string{
102			`_pgs.example.com: "antonio-blgo" is invalid: user "antonio" has no project "blgo"`,
103			`_prose.example.com: "nobody" is invalid: user "nobody" not found`,
104			`_sish.example.com: "SHA256:someone-else" is invalid: no record is your username or the SHA256 fingerprint of one of your keys`,
105		}},
106		{"tuns record with whitespace", "", map[string][]string{"sish": {fingerprint + " "}}, []string{
107			fmt.Sprintf(`_sish.example.com: %q is invalid: it has surrounding whitespace, and tuns needs an exact match`, fingerprint+" "),
108		}},
109		{"tuns username", "", map[string][]string{"tuns": {"antonio"}}, []string{
110			`_tuns.example.com: "antonio" is valid, tunnels from any of your keys can use this domain`,
111			"_sish.example.com: no TXT record",
112		}},
113		{"wildcard", "*.example.com", map[string][]string{"pgs": {"antonio-blog"}, "sish": {"antonio"}}, []string{
114			"_pgs.*.example.com: wildcard domains are not supported\r\n",
115			"_prose.*.example.com: wildcard domains are not supported\r\n",
116			`.example.com: "antonio" is valid, tunnels from any of your keys can use this domain`,
117			"_tuns.",
118		}},
119	}
120	for _, tt := range tests {
121		t.Run(tt.name, func(t *testing.T) {
122			session := &fakeSession{}
123			cmd := &Cmd{
124				User:    me,
125				Session: session,
126				Log:     slog.New(slog.DiscardHandler),
127				Dbpool:  dbpool,
128				lookupHost: func(host string) (string, []string, error) {
129					return "", nil, &net.DNSError{IsNotFound: true}
130				},
131				lookupTXT: func(host, space string) ([]string, error) {
132					if tt.domain == "" && host != "example.com" {
133						t.Fatalf("looked up %q", host)
134					}
135					if tt.domain != "" && (!strings.HasSuffix(host, ".example.com") || strings.Contains(host, "*")) {
136						t.Fatalf("looked up %q for %q", host, tt.domain)
137					}
138					if v, ok := tt.records[space]; ok {
139						return v, nil
140					}
141					return nil, router.ErrNoTXTRecord
142				},
143			}
144			domain := tt.domain
145			if domain == "" {
146				domain = " Example.COM. "
147			}
148			if err := cmd.dnsCheck(domain); err != nil {
149				t.Fatal(err)
150			}
151			out := session.String()
152			for _, line := range tt.want {
153				if !strings.Contains(out, line) {
154					t.Errorf("missing %q in output:\n%s", line, out)
155				}
156			}
157		})
158	}
159
160	pointers := []struct {
161		name  string
162		hosts map[string][]string
163		cname string
164		want  string
165	}{
166		{"none", map[string][]string{}, "", "example.com: no CNAME or A record"},
167		{"cname to a region", map[string][]string{
168			"example.com": {"10.0.0.3"},
169			"pgs.sh":      {"10.0.0.1"},
170			"prose.sh":    {"10.0.0.2"},
171			"tuns.sh":     {"10.0.0.3"},
172			"nue.tuns.sh": {"10.0.0.4"},
173		}, "tuns.pico.sh", "example.com: CNAME tuns.pico.sh, resolves to 10.0.0.3, which is tuns.sh"},
174		{"apex A", map[string][]string{
175			"example.com": {"10.0.0.9", "10.0.0.1"},
176			"pgs.sh":      {"10.0.0.1"},
177		}, "", "example.com: resolves to 10.0.0.9, 10.0.0.1, which is pgs.sh"},
178		{"elsewhere", map[string][]string{
179			"example.com": {"10.0.0.9"},
180			"pgs.sh":      {"10.0.0.1"},
181		}, "", "example.com: resolves to 10.0.0.9, which is not pgs.sh, prose.sh, tuns.sh, nue.tuns.sh"},
182	}
183	for _, tt := range pointers {
184		t.Run("pointer "+tt.name, func(t *testing.T) {
185			session := &fakeSession{}
186			cmd := &Cmd{
187				User:          me,
188				Session:       session,
189				Dbpool:        dbpool,
190				DomainTargets: []string{"pgs.sh", "prose.sh", "tuns.sh", "nue.tuns.sh"},
191				lookupHost: func(host string) (string, []string, error) {
192					addrs, ok := tt.hosts[host]
193					if !ok {
194						return "", nil, &net.DNSError{IsNotFound: true}
195					}
196					if host == "example.com" {
197						return tt.cname, addrs, nil
198					}
199					return "", addrs, nil
200				},
201				lookupTXT: func(host, space string) ([]string, error) {
202					return nil, router.ErrNoTXTRecord
203				},
204			}
205			if err := cmd.dnsCheck("example.com"); err != nil {
206				t.Fatal(err)
207			}
208			if out := session.String(); !strings.Contains(out, tt.want+"\r\n") {
209				t.Errorf("missing %q in output:\n%s", tt.want, out)
210			}
211		})
212	}
213
214	cmd := &Cmd{User: me, Session: &fakeSession{}, Dbpool: dbpool}
215	for _, bad := range []string{"https://example.com/", "*.", "a.*.example.com"} {
216		if err := cmd.dnsCheck(bad); err == nil {
217			t.Errorf("%q was accepted as a domain", bad)
218		}
219	}
220}