1
package client
2
 
3
import (
4
	"context"
5
	"crypto/tls"
6
	"crypto/x509"
7
	"errors"
8
	"fmt"
9
	"time"
10
 
11
	pkiImpl "go.miragespace.co/specter/pki"
12
	"go.miragespace.co/specter/spec/pki"
13
	"go.miragespace.co/specter/spec/protocol"
14
	"go.miragespace.co/specter/spec/rpc"
15
	"go.miragespace.co/specter/spec/transport"
16
 
17
	"go.uber.org/zap"
18
)
19
 
20
func (c *Client) updateTransportCert() error {
21
	cert, err := tls.X509KeyPair([]byte(c.Configuration.Certificate), []byte(c.Configuration.PrivKey))
22
	if err != nil {
23
		return fmt.Errorf("error parsing certificate: %w", err)
24
	}
25
	if tp, ok := c.ServerTransport.(transport.ClientTransport); ok {
26
		if err := tp.WithClientCertificate(cert); err != nil {
27
			return err
28
		}
29
	} else {
30
		return fmt.Errorf("transport does not support client certificate override")
31
	}
32
	return nil
33
}
34
 
35
func (c *Client) Register(ctx context.Context) error {
36
	c.configMu.Lock()
37
	defer c.configMu.Unlock()
38
 
39
	if c.Configuration.Certificate != "" {
40
		clientCert, err := tls.X509KeyPair([]byte(c.Configuration.Certificate), []byte(c.Configuration.PrivKey))
41
		if err != nil {
42
			return fmt.Errorf("error parsing certificate: %w", err)
43
		}
44
		cert, err := x509.ParseCertificate(clientCert.Certificate[0])
45
		if err != nil {
46
			return fmt.Errorf("error parsing certificate: %w", err)
47
		}
48
 
49
		// Check if certificate needs renewal
50
		if shouldRenewCertificate(cert) && c.PKIClient != nil {
51
			c.Logger.Info("Certificate is within renewal window, attempting renewal")
52
			if err := c.renewCertificate(ctx, cert); err != nil {
53
				c.Logger.Warn("Failed to renew certificate, continuing with existing certificate", zap.Error(err))
54
			} else {
55
				// Re-parse the updated certificate after successful renewal
56
				clientCert, err = tls.X509KeyPair([]byte(c.Configuration.Certificate), []byte(c.Configuration.PrivKey))
57
				if err != nil {
58
					return fmt.Errorf("error parsing renewed certificate: %w", err)
59
				}
60
				cert, err = x509.ParseCertificate(clientCert.Certificate[0])
61
				if err != nil {
62
					return fmt.Errorf("error parsing renewed certificate: %w", err)
63
				}
64
			}
65
		}
66
 
67
		resp, err := retryRPC(c, ctx, func(node *protocol.Node) (*protocol.ClientPingResponse, error) {
68
			ctx = rpc.WithNode(ctx, node)
69
			return c.tunnelClient.Ping(ctx, &protocol.ClientPingRequest{})
70
		})
71
		if err != nil {
72
			return err
73
		}
74
		root := resp.GetApex()
75
		c.rootDomain.Store(root)
76
 
77
		identity, err := pki.ExtractCertificateIdentity(cert)
78
		if err != nil {
79
			return fmt.Errorf("failed to extract certificate identity: %w", err)
80
		}
81
		c.Logger.Info("Reusing existing client certificate", zap.Object("identity", identity))
82
 
83
		return nil
84
	}
85
 
86
	if c.PKIClient == nil {
87
		return errors.New("no client certificate found: please ensure your client is registered with apex first with the tunnel subcommand")
88
	}
89
 
90
	c.Logger.Info("Obtaining a new client certificate from apex")
91
 
92
	privKey, err := pki.UnmarshalPrivateKey([]byte(c.Configuration.PrivKey))
93
	if err != nil {
94
		return fmt.Errorf("failed to parse private key: %w", err)
95
	}
96
 
97
	pkiReq, err := pkiImpl.CreateRequest(privKey)
98
	if err != nil {
99
		return fmt.Errorf("failed to create certificate request: %w", err)
100
	}
101
 
102
	pkiResp, err := c.PKIClient.RequestCertificate(c.parentCtx, pkiReq)
103
	if err != nil {
104
		return fmt.Errorf("failed to obtain a client certificate: %w", err)
105
	}
106
 
107
	cert, err := x509.ParseCertificate(pkiResp.GetCertDer())
108
	if err != nil {
109
		return fmt.Errorf("invalid certificate from PKIService: %w", err)
110
	}
111
 
112
	c.Configuration.Certificate = string(pkiResp.GetCertPem())
113
 
114
	identity, err := pki.ExtractCertificateIdentity(cert)
115
	if err != nil {
116
		return fmt.Errorf("failed to extract certificate identity: %w", err)
117
	}
118
	c.Logger.Info("Client certificate obtained", zap.Object("identity", identity))
119
 
120
	if err := c.updateTransportCert(); err != nil {
121
		return fmt.Errorf("failed to update transport certificate: %w", err)
122
	}
123
	// Registration runs before forwarding starts; renewal keeps these loggers stable.
124
	c.Logger = c.Logger.With(zap.Uint64("id", c.ServerTransport.Identity().GetId()))
125
	c.forwarder.logger = c.Logger
126
 
127
	if err := c.bootstrap(c.parentCtx, c.Configuration.Apex); err != nil {
128
		return fmt.Errorf("failed to bootstrap with certificate: %w", err)
129
	}
130
 
131
	resp, err := retryRPC(c, ctx, func(node *protocol.Node) (*protocol.RegisterIdentityResponse, error) {
132
		ctx = rpc.WithNode(ctx, node)
133
		return c.tunnelClient.RegisterIdentity(ctx, &protocol.RegisterIdentityRequest{})
134
	})
135
	if err != nil {
136
		return err
137
	}
138
 
139
	root := resp.GetApex()
140
	c.rootDomain.Store(root)
141
 
142
	if err := c.Configuration.writeFile(); err != nil {
143
		c.Logger.Error("Error saving token to config file", zap.Error(err))
144
	}
145
 
146
	return nil
147
}
148
 
149
// shouldRenewCertificate checks if the certificate is within the renewal window.
150
func shouldRenewCertificate(cert *x509.Certificate) bool {
151
	return time.Until(cert.NotAfter) <= renewalWindow
152
}
153
 
154
// renewCertificate performs certificate renewal using the PKI service.
155
// It must be called with configMu held (or by a caller that holds configMu).
156
func (c *Client) renewCertificate(ctx context.Context, oldCert *x509.Certificate) error {
157
	privKeyPEM := c.Configuration.PrivKey
158
 
159
	resp, newCert, err := c.performRenewalRPC(ctx, oldCert, privKeyPEM)
160
	if err != nil {
161
		return err
162
	}
163
 
164
	return c.updateConfigurationWithCert(resp, newCert)
165
}
166
 
167
// performRenewalRPC performs the network call to renew the certificate.
168
// It does not hold configMu.
169
func (c *Client) performRenewalRPC(ctx context.Context, oldCert *x509.Certificate, privKeyPEM string) (*protocol.CertificateResponse, *x509.Certificate, error) {
170
	if c.PKIClient == nil {
171
		return nil, nil, errors.New("PKI client not available for certificate renewal")
172
	}
173
 
174
	c.Logger.Info("Renewing client certificate",
175
		zap.Time("notAfter", oldCert.NotAfter),
176
		zap.Duration("timeUntilExpiry", time.Until(oldCert.NotAfter)))
177
 
178
	privKey, err := pki.UnmarshalPrivateKey([]byte(privKeyPEM))
179
	if err != nil {
180
		return nil, nil, fmt.Errorf("failed to parse private key: %w", err)
181
	}
182
 
183
	renewReq, err := pkiImpl.CreateRenewalRequest(privKey, oldCert.Raw)
184
	if err != nil {
185
		return nil, nil, fmt.Errorf("failed to create renewal request: %w", err)
186
	}
187
 
188
	renewResp, err := c.PKIClient.RenewCertificate(ctx, renewReq)
189
	if err != nil {
190
		return nil, nil, fmt.Errorf("failed to renew certificate: %w", err)
191
	}
192
 
193
	newCert, err := x509.ParseCertificate(renewResp.GetCertDer())
194
	if err != nil {
195
		return nil, nil, fmt.Errorf("invalid renewed certificate from PKIService: %w", err)
196
	}
197
 
198
	return renewResp, newCert, nil
199
}
200
 
201
// updateConfigurationWithCert updates the client configuration with the renewed certificate.
202
// It must be called with configMu held.
203
func (c *Client) updateConfigurationWithCert(renewResp *protocol.CertificateResponse, newCert *x509.Certificate) error {
204
	c.Configuration.Certificate = string(renewResp.GetCertPem())
205
 
206
	identity, err := pki.ExtractCertificateIdentity(newCert)
207
	if err != nil {
208
		return fmt.Errorf("failed to extract identity from renewed certificate: %w", err)
209
	}
210
	c.Logger.Info("Certificate renewed successfully",
211
		zap.Object("identity", identity),
212
		zap.Time("newNotAfter", newCert.NotAfter))
213
 
214
	if err := c.updateTransportCert(); err != nil {
215
		return fmt.Errorf("failed to update transport certificate: %w", err)
216
	}
217
 
218
	if err := c.Configuration.writeFile(); err != nil {
219
		c.Logger.Error("Error saving renewed certificate to config file", zap.Error(err))
220
	}
221
 
222
	return nil
223
}
224
 
225
// certificateMaintainer runs in the background to periodically check and renew
226
// the client certificate before it expires. This is important for long-running clients.
227
func (c *Client) certificateMaintainer(ctx context.Context) {
228
	defer c.closeWg.Done()
229
 
230
	// Don't run if PKI client is not available
231
	if c.PKIClient == nil {
232
		return
233
	}
234
 
235
	ticker := time.NewTicker(certCheckInterval)
236
	defer ticker.Stop()
237
 
238
	for {
239
		select {
240
		case <-c.closeCh:
241
			return
242
		case <-ctx.Done():
243
			return
244
		case <-ticker.C:
245
			c.checkAndRenewCertificate(ctx)
246
		}
247
	}
248
}
249
 
250
// checkAndRenewCertificate checks if the certificate needs renewal and performs it if necessary.
251
func (c *Client) checkAndRenewCertificate(ctx context.Context) {
252
	c.configMu.RLock()
253
	certPEM := c.Configuration.Certificate
254
	privKeyPEM := c.Configuration.PrivKey
255
	c.configMu.RUnlock()
256
 
257
	if certPEM == "" {
258
		return
259
	}
260
 
261
	// Parse the certificate to check expiry
262
	clientCert, err := tls.X509KeyPair([]byte(certPEM), []byte(privKeyPEM))
263
	if err != nil {
264
		c.Logger.Error("Failed to parse certificate during maintenance check", zap.Error(err))
265
		return
266
	}
267
	cert, err := x509.ParseCertificate(clientCert.Certificate[0])
268
	if err != nil {
269
		c.Logger.Error("Failed to parse certificate during maintenance check", zap.Error(err))
270
		return
271
	}
272
 
273
	if !shouldRenewCertificate(cert) {
274
		c.Logger.Debug("Certificate does not need renewal yet",
275
			zap.Time("notAfter", cert.NotAfter),
276
			zap.Duration("timeUntilExpiry", time.Until(cert.NotAfter)))
277
		return
278
	}
279
 
280
	c.Logger.Info("Background certificate maintenance: renewing certificate")
281
 
282
	// Perform RPC without holding the lock
283
	renewResp, newCert, err := c.performRenewalRPC(ctx, cert, privKeyPEM)
284
	if err != nil {
285
		c.Logger.Error("Background certificate renewal failed", zap.Error(err))
286
		return
287
	}
288
 
289
	// Upgrade to write lock for updating configuration
290
	c.configMu.Lock()
291
	defer c.configMu.Unlock()
292
 
293
	// Check if configuration has changed while we were renewing
294
	if c.Configuration.Certificate != certPEM {
295
		c.Logger.Warn("Configuration changed during certificate renewal, aborting update")
296
		return
297
	}
298
 
299
	if err := c.updateConfigurationWithCert(renewResp, newCert); err != nil {
300
		c.Logger.Error("Failed to update configuration with renewed certificate", zap.Error(err))
301
	}
302
}