From aa784a9f15e79a3011f25c081fb87f52326a8ad3 Mon Sep 17 00:00:00 2001
From: Steffen Vogel <post@steffenvogel.de>
Date: Sun, 21 Aug 2022 10:54:24 +0200
Subject: [PATCH] device: add new handshake handler and keylog writer

This change adds support for a new environment variable 'WG_KEYLOGFILE'
in resemblance to the 'SSLKEYLOGFILE' environment variable used by
curl, Chrome & Firefox to log ephemeral TLS encryption keys

When set, wireguard-go will log ephemeral keys generated during
each handshake to a file specified by the environment variable in the
WireGuard key log format.

The format used is the same as then one generated by the
extract-handshakes.sh script.

See also:
- https://git.zx2c4.com/wireguard-tools/tree/contrib/extract-handshakes
- https://wiki.wireshark.org/WireGuard#key-log-format
- https://everything.curl.dev/usingcurl/tls/sslkeylogfile

Signed-off-by: Steffen Vogel <post@steffenvogel.de>
---
 device/device.go         | 37 +++++++++++++++++++++++++++++++++----
 device/noise-protocol.go |  9 +++++++++
 device/noise-types.go    |  9 +++++++++
 main.go                  | 26 ++++++++++++++++++++++++++
 tai64n/tai64n.go         |  6 +++++-
 5 files changed, 82 insertions(+), 5 deletions(-)

diff --git a/device/device.go b/device/device.go
index 3625608..8721c49 100644
--- a/device/device.go
+++ b/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")
+	})
+}
diff --git a/device/noise-protocol.go b/device/noise-protocol.go
index ffa0452..dd85bbf 100644
--- a/device/noise-protocol.go
+++ b/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
 }
 
diff --git a/device/noise-types.go b/device/noise-types.go
index 6e850e7..d17bb9f 100644
--- a/device/noise-types.go
+++ b/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)
+}
diff --git a/main.go b/main.go
index b35ac29..769cba5 100644
--- a/main.go
+++ b/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)
diff --git a/tai64n/tai64n.go b/tai64n/tai64n.go
index de0f5bf..9484c26 100644
--- a/tai64n/tai64n.go
+++ b/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])))
 }
-- 
2.36.0

