Add new handshake handler and keylog writer

device/device.go

@@ -6,6 +6,9 @@
 package device
 
 import (
+	"encoding/base64"
+	"fmt"
+	"io"
 	"runtime"
 	"sync"
 	"sync/atomic"
@@ -17,6 +20,8 @@ import (
 	"golang.zx2c4.com/wireguard/tun"
 )
 
+type HandshakeHandler func(t time.Time, ls NoisePrivateKey, rs NoisePublicKey, le NoisePrivateKey, ps NoisePresharedKey)
+
 type Device struct {
 	state struct {
 		// state holds the device's state. It is accessed atomically.
@@ -85,6 +90,8 @@ type Device struct {
 		mtu    int32
 	}
 
+	keyLogHandler HandshakeHandler
+
 	ipcMutex sync.RWMutex
 	closed   chan struct{}
 	log      *Logger
@@ -94,10 +101,9 @@ type Device struct {
 // There are three states: down, up, closed.
 // Transitions:
 //
-//   down -----+
-//     ↑↓      ↓
-//     up -> closed
-//
+//	down -----+
+//	  ↑↓      ↓
+//	  up -> closed
 type deviceState uint32
 
 //go:generate go run golang.org/x/tools/cmd/stringer -type deviceState -trimprefix=deviceState
@@ -523,3 +529,26 @@ func (device *Device) BindClose() error {
 	device.net.Unlock()
 	return err
 }
+
+func (device *Device) OnHandshake(hdlr HandshakeHandler) {
+	device.keyLogHandler = hdlr
+}
+
+func (device *Device) WriteKeyLog(wr io.Writer) {
+	mu := sync.Mutex{}
+
+	device.OnHandshake(func(t time.Time, ls NoisePrivateKey, rs NoisePublicKey, le NoisePrivateKey, ps NoisePresharedKey) {
+		mu.Lock()
+		defer mu.Unlock()
+
+		fmt.Fprintf(wr, "LOCAL_STATIC_PRIVATE_KEY=%s\n", base64.StdEncoding.EncodeToString(ls[:]))
+		fmt.Fprintf(wr, "REMOTE_STATIC_PUBLIC_KEY=%s\n", base64.StdEncoding.EncodeToString(rs[:]))
+		fmt.Fprintf(wr, "LOCAL_EPHEMERAL_PRIVATE_KEY=%s\n", base64.StdEncoding.EncodeToString(le[:]))
+
+		if !ps.IsZero() {
+			fmt.Fprintf(wr, "PRESHARED_KEY=%s\n", base64.StdEncoding.EncodeToString(ps[:]))
+		}
+
+		device.log.Verbosef("Writing new ephemeral key to keylog")
+	})
+}

device/noise-protocol.go

@@ -243,6 +243,11 @@ func (device *Device) CreateMessageInitiation(peer *Peer) (*MessageInitiation, e
 
 	handshake.mixHash(msg.Timestamp[:])
 	handshake.state = handshakeInitiationCreated
+
+	if device.keyLogHandler != nil {
+		go device.keyLogHandler(handshake.lastTimestamp.Time(), device.staticIdentity.privateKey, handshake.remoteStatic, handshake.localEphemeral, handshake.presharedKey)
+	}
+
 	return &msg, nil
 }
 
@@ -414,6 +419,10 @@ func (device *Device) CreateMessageResponse(peer *Peer) (*MessageResponse, error
 
 	handshake.state = handshakeResponseCreated
 
+	if device.keyLogHandler != nil {
+		go device.keyLogHandler(handshake.lastTimestamp.Time(), device.staticIdentity.privateKey, handshake.remoteStatic, handshake.localEphemeral, handshake.presharedKey)
+	}
+
 	return &msg, nil
 }
 

device/noise-types.go

@@ -76,3 +76,12 @@ func (key NoisePublicKey) Equals(tar NoisePublicKey) bool {
 func (key *NoisePresharedKey) FromHex(src string) error {
 	return loadExactHex(key[:], src)
 }
+
+func (key NoisePresharedKey) Equals(tar NoisePresharedKey) bool {
+	return subtle.ConstantTimeCompare(key[:], tar[:]) == 1
+}
+
+func (key NoisePresharedKey) IsZero() bool {
+	var zero NoisePresharedKey
+	return key.Equals(zero)
+}

main.go

@@ -30,6 +30,7 @@ const (
 	ENV_WG_TUN_FD             = "WG_TUN_FD"
 	ENV_WG_UAPI_FD            = "WG_UAPI_FD"
 	ENV_WG_PROCESS_FOREGROUND = "WG_PROCESS_FOREGROUND"
+	ENV_WG_KEYLOG             = "WG_KEYLOGFILE"
 )
 
 func printUsage() {
@@ -152,6 +153,22 @@ func main() {
 		os.Exit(ExitSetupFailed)
 	}
 
+	// open keylog file
+
+	keyLog, err := func() (*os.File, error) {
+		fn := os.Getenv(ENV_WG_KEYLOG)
+		if fn == "" {
+			return nil, nil
+		}
+
+		return os.OpenFile(fn, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0600)
+	}()
+	if err != nil {
+		logger.Errorf("failed to open keylog error: %v", err)
+		os.Exit(ExitSetupFailed)
+		return
+	}
+
 	// open UAPI file (or use supplied fd)
 
 	fileUAPI, err := func() (*os.File, error) {
@@ -174,6 +191,7 @@ func main() {
 		os.Exit(ExitSetupFailed)
 		return
 	}
+
 	// daemonize the process
 
 	if !foreground {
@@ -248,6 +266,14 @@ func main() {
 
 	logger.Verbosef("UAPI listener started")
 
+	// start writing handshakes to keylog file
+
+	if keyLog != nil {
+		device.WriteKeyLog(keyLog)
+
+		logger.Verbosef("Keylog writer started")
+	}
+
 	// wait for program to terminate
 
 	signal.Notify(term, syscall.SIGTERM)

tai64n/tai64n.go

@@ -37,5 +37,9 @@ func (t1 Timestamp) After(t2 Timestamp) bool {
 }
 
 func (t Timestamp) String() string {
-	return time.Unix(int64(binary.BigEndian.Uint64(t[:8])-base), int64(binary.BigEndian.Uint32(t[8:12]))).String()
+	return t.Time().String()
+}
+
+func (t Timestamp) Time() time.Time {
+	return time.Unix(int64(binary.BigEndian.Uint64(t[:8])-base), int64(binary.BigEndian.Uint32(t[8:12])))
 }