1
package client
2
 
3
import (
4
	"bytes"
5
	"crypto"
6
	"crypto/rand"
7
	"crypto/rsa"
8
	"crypto/sha256"
9
	"crypto/tls"
10
	"crypto/x509"
11
	"crypto/x509/pkix"
12
	"math/big"
13
	"net"
14
	"os"
15
	"testing"
16
	"time"
17
 
18
	"go.miragespace.co/specter/spec/chord"
19
	"go.miragespace.co/specter/spec/mocks"
20
	"go.miragespace.co/specter/spec/pki"
21
	"go.miragespace.co/specter/spec/protocol"
22
 
23
	"github.com/stretchr/testify/mock"
24
	"github.com/stretchr/testify/require"
25
	"github.com/zhangyunhao116/skipmap"
26
	"go.uber.org/zap/zaptest"
27
)
28
 
29
func TestKeylessSign(t *testing.T) {
30
	as := require.New(t)
31
	logger := zaptest.NewLogger(t)
32
 
33
	file, err := os.CreateTemp("", "client")
34
	as.NoError(err)
35
	defer os.Remove(file.Name())
36
 
37
	ctx := t.Context()
38
 
39
	token := &protocol.ClientToken{
40
		Token: []byte("test"),
41
	}
42
	cl := &protocol.Node{
43
		Id: chord.Random(),
44
	}
45
 
46
	hostname := "custom.domain.com"
47
 
48
	der, cert, key := makeCertificate(as, logger, cl, token, nil)
49
	privKey, err := pki.UnmarshalPrivateKey([]byte(key))
50
	as.NoError(err)
51
 
52
	cfg := &Config{
53
		path:        file.Name(),
54
		router:      skipmap.NewString[route](),
55
		Apex:        testApex,
56
		Certificate: cert,
57
		PrivKey:     key,
58
		Tunnels: []Tunnel{
59
			{
60
				Target: "tcp://127.0.0.1:2345",
61
			},
62
		},
63
	}
64
	as.NoError(cfg.validate())
65
 
66
	_, providerX509Cert, providerPrivateKey := testMakeRSACert(as)
67
	providerCert := &tls.Certificate{
68
		Certificate: [][]byte{providerX509Cert.Raw},
69
		PrivateKey:  providerPrivateKey,
70
		Leaf:        providerX509Cert,
71
	}
72
 
73
	rawMsg := []byte("hello world")
74
	digest := sha256.Sum256(rawMsg)
75
	algoUsed := protocol.KeylessSignRequest_SHA256
76
	providerSig, err := providerPrivateKey.Sign(rand.Reader, digest[:], &rsa.PSSOptions{
77
		SaltLength: rsa.PSSSaltLengthAuto,
78
		Hash:       crypto.SHA256,
79
	})
80
	as.NoError(err)
81
 
82
	m := func(s *mocks.TunnelService, t1 *mocks.MemoryTransport, publishCall *mock.Call) {
83
		s.Keyless.On("GetCertificate", mock.Anything, mock.MatchedBy(func(req *protocol.KeylessGetCertificateRequest) bool {
84
			return powValidateFunc(hostname, privKey)(req.GetProof(), req.GetHostname())
85
		})).Return(&protocol.KeylessGetCertificateResponse{
86
			Certificates: providerCert.Certificate,
87
		}, nil)
88
		s.Keyless.On("Sign", mock.Anything, mock.MatchedBy(func(req *protocol.KeylessSignRequest) bool {
89
			powValid := powValidateFunc(hostname, privKey)(req.GetProof(), req.GetHostname())
90
			if !powValid {
91
				return false
92
			}
93
			return req.GetAlgo() == algoUsed && bytes.Equal(digest[:], req.GetDigest())
94
		})).Return(&protocol.KeylessSignResponse{
95
			Signature: providerSig,
96
		}, nil)
97
 
98
		defaultNoHostnames(s)
99
		transportHelper(t1, der)
100
	}
101
 
102
	listenCfg := &net.ListenConfig{}
103
	sListener, err := listenCfg.Listen(ctx, "tcp", "127.0.0.1:0")
104
	as.NoError(err)
105
	defer sListener.Close()
106
 
107
	client, _, assertion := setupClient(t, as, ctx, logger, nil, cfg, nil, m, false, 1)
108
	defer assertion()
109
	defer client.Close()
110
 
111
	client.ServerListener = sListener
112
 
113
	client.Start(ctx)
114
 
115
	keylessCert, err := client.getCertificate(&tls.ClientHelloInfo{
116
		ServerName: hostname,
117
	})
118
	as.NoError(err)
119
	as.NotNil(keylessCert)
120
 
121
	keylessSigner, ok := keylessCert.PrivateKey.(crypto.Signer)
122
	as.True(ok)
123
 
124
	keylessSignature, err := keylessSigner.Sign(rand.Reader, digest[:], &rsa.PSSOptions{
125
		SaltLength: rsa.PSSSaltLengthAuto,
126
		Hash:       crypto.SHA256,
127
	})
128
	as.NoError(err)
129
	as.NotNil(keylessSignature)
130
	as.EqualValues(providerSig, keylessSignature)
131
}
132
 
133
func testMakeRSACert(as *require.Assertions) (derBytes []byte, x509Cert *x509.Certificate, privateKey *rsa.PrivateKey) {
134
	var err error
135
 
136
	privateKey, err = rsa.GenerateKey(rand.Reader, 2048)
137
	as.NoError(err)
138
 
139
	template := x509.Certificate{
140
		SerialNumber: big.NewInt(1),
141
		Subject: pkix.Name{
142
			Organization: []string{"dev"},
143
		},
144
		NotBefore: time.Now(),
145
		NotAfter:  time.Now().Add(time.Hour * 24 * 365), // Valid for 1 year
146
 
147
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageKeyEncipherment,
148
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
149
		BasicConstraintsValid: true,
150
	}
151
 
152
	derBytes, err = x509.CreateCertificate(rand.Reader, &template, &template, &privateKey.PublicKey, privateKey)
153
	as.NoError(err)
154
 
155
	x509Cert, err = x509.ParseCertificate(derBytes)
156
	as.NoError(err)
157
 
158
	return
159
}