diff --git a/main.go b/main.go index 61e7e33..4708f74 100644 --- a/main.go +++ b/main.go @@ -597,6 +597,12 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey( return nil }) + client.OnTokenUpdate(func(token string) { + if wgService != nil { + wgService.SetToken(token) + } + }) + // Connect to the WebSocket server if err := client.Connect(); err != nil { logger.Fatal("Failed to connect to server: %v", err) diff --git a/websocket/client.go b/websocket/client.go index 2706eee..98c9388 100644 --- a/websocket/client.go +++ b/websocket/client.go @@ -27,7 +27,8 @@ type Client struct { isConnected bool reconnectMux sync.RWMutex - onConnect func() error + onConnect func() error + onTokenUpdate func(token string) } type ClientOption func(*Client) @@ -45,6 +46,10 @@ func (c *Client) OnConnect(callback func() error) { c.onConnect = callback } +func (c *Client) OnTokenUpdate(callback func(token string)) { + c.onTokenUpdate = callback +} + // NewClient creates a new Newt client func NewClient(newtID, secret string, endpoint string, opts ...ClientOption) (*Client, error) { config := &Config{ @@ -270,6 +275,8 @@ func (c *Client) establishConnection() error { return fmt.Errorf("failed to get token: %w", err) } + c.onTokenUpdate(token) + // Parse the base URL to determine protocol and hostname baseURL, err := url.Parse(c.baseURL) if err != nil { diff --git a/wg/wg.go b/wg/wg.go index a3f7f9e..7330fdf 100644 --- a/wg/wg.go +++ b/wg/wg.go @@ -60,6 +60,7 @@ type WireGuardService struct { stopHolepunch chan struct{} host string serverPubKey string + token string } // Add this type definition @@ -181,6 +182,10 @@ func (s *WireGuardService) SetServerPubKey(serverPubKey string) { s.serverPubKey = serverPubKey } +func (s *WireGuardService) SetToken(token string) { + s.token = token +} + func (s *WireGuardService) LoadRemoteConfig() error { err := s.client.SendMessage("newt/wg/get-config", map[string]interface{}{ @@ -624,6 +629,11 @@ func (s *WireGuardService) reportPeerBandwidth() error { } func (s *WireGuardService) sendUDPHolePunch(serverAddr string) error { + + if s.serverPubKey == "" || s.token == "" { + return fmt.Errorf("server public key or token is not set") + } + // Parse server address serverSplit := strings.Split(serverAddr, ":") if len(serverSplit) < 2 { @@ -665,8 +675,10 @@ func (s *WireGuardService) sendUDPHolePunch(serverAddr string) error { // Create JSON payload payload := struct { NewtID string `json:"newtId"` + Token string `json:"token"` }{ NewtID: s.newtId, + Token: s.token, } // Convert payload to JSON @@ -690,7 +702,6 @@ func (s *WireGuardService) sendUDPHolePunch(serverAddr string) error { return nil } -// Add a new function to encrypt the payload func (s *WireGuardService) encryptPayload(payload []byte) (interface{}, error) { // Generate an ephemeral keypair for this message ephemeralPrivateKey, err := wgtypes.GeneratePrivateKey() @@ -705,18 +716,18 @@ func (s *WireGuardService) encryptPayload(payload []byte) (interface{}, error) { return nil, fmt.Errorf("failed to parse server public key: %v", err) } - // Perform Diffie-Hellman key exchange - var serverPubKeyFixed [32]byte - copy(serverPubKeyFixed[:], serverPubKey[:]) - + // Use X25519 for key exchange (replacing deprecated ScalarMult) var ephPrivKeyFixed [32]byte copy(ephPrivKeyFixed[:], ephemeralPrivateKey[:]) - var sharedSecret [32]byte - curve25519.ScalarMult(&sharedSecret, &ephPrivKeyFixed, &serverPubKeyFixed) + // Perform X25519 key exchange + sharedSecret, err := curve25519.X25519(ephPrivKeyFixed[:], serverPubKey[:]) + if err != nil { + return nil, fmt.Errorf("failed to perform X25519 key exchange: %v", err) + } // Create an AEAD cipher using the shared secret - aead, err := chacha20poly1305.New(sharedSecret[:]) + aead, err := chacha20poly1305.New(sharedSecret) if err != nil { return nil, fmt.Errorf("failed to create AEAD cipher: %v", err) }