Add proxy IP to audit log

Author: tsipinakis

auditlog/message/connect.go

@@ -2,8 +2,9 @@ package message
 
 // PayloadConnect is the payload for TypeConnect messages.
 type PayloadConnect struct {
-	RemoteAddr string `json:"remoteAddr" yaml:"remoteAddr"` // RemoteAddr contains the IP address of the connecting user.
-	Country    string `json:"country" yaml:"country"`       // Country contains the country code looked up from the IP address. Contains "XX" if the lookup failed.
+	RemoteAddr string `json:"remoteAddr" yaml:"remoteAddr"`          // RemoteAddr contains the IP address of the connecting user.
+	ProxyAddr  string `json:"proxyAddr,omitempty"  yaml:"proxyAddr"` // ProxyAddr contains the IP adress of the proxy used (if behind a load balancer)
+	Country    string `json:"country" yaml:"country"`                // Country contains the country code looked up from the IP address. Contains "XX" if the lookup failed.
 }
 
 // Equals compares two PayloadConnect datasets.

internal/auditlog/codec/asciinema/encoder.go

@@ -62,6 +62,7 @@ func (e *encoder) Encode(messages <-chan message.Message, storage storage.Writer
 	startTime := int64(0)
 	headerWritten := false
 	var ip = ""
+	var proxy *string
 	var username *string
 	const shell = "/bin/sh"
 	for {
@@ -70,11 +71,12 @@ func (e *encoder) Encode(messages <-chan message.Message, storage storage.Writer
 			break
 		}
 		var err error
-		startTime, headerWritten, ip, username, err = e.encodeMessage(
+		startTime, headerWritten, ip, proxy, username, err = e.encodeMessage(
 			startTime,
 			msg,
 			&asciicastHeader,
 			ip,
+			proxy,
 			storage,
 			username,
 			headerWritten,
@@ -106,11 +108,12 @@ func (e *encoder) encodeMessage(
 	msg message.Message,
 	asciicastHeader *Header,
 	ip string,
+	proxy *string,
 	storage storage.Writer,
 	username *string,
 	headerWritten bool,
 	shell string,
-) (int64, bool, string, *string, error) {
+) (int64, bool, string, *string, *string, error) {
 	if msg.MessageType == message.TypeConnect {
 		startTime = msg.Timestamp
 		asciicastHeader.Timestamp = int(startTime / 1000000000)
@@ -121,11 +124,11 @@ func (e *encoder) encodeMessage(
 	case message.TypeConnect:
 		ip, username = e.handleConnect(storage, msg, startTime, country, username)
 	case message.TypeAuthPasswordSuccessful:
-		ip, username = e.handleAuthPasswordSuccessful(storage, msg, startTime, ip, country)
+		ip, username = e.handleAuthPasswordSuccessful(storage, msg, startTime, ip, proxy, country)
 	case message.TypeAuthPubKeySuccessful:
-		ip, username = e.handleAuthPubkeySuccessful(storage, msg, startTime, ip, country)
+		ip, username = e.handleAuthPubkeySuccessful(storage, msg, startTime, ip, proxy, country)
 	case message.TypeHandshakeSuccessful:
-		ip, username = e.handleHandshakeSuccessful(storage, msg, startTime, ip, country)
+		ip, username = e.handleHandshakeSuccessful(storage, msg, startTime, ip, proxy, country)
 	case message.TypeChannelRequestSetEnv:
 		payload := msg.Payload.(message.PayloadChannelRequestSetEnv)
 		asciicastHeader.Env[payload.Name] = payload.Value
@@ -142,36 +145,40 @@ func (e *encoder) encodeMessage(
 		startTime, headerWritten, err = e.handleIO(startTime, msg, asciicastHeader, headerWritten, shell, storage)
 	}
 	if err != nil {
-		return startTime, headerWritten, ip, username, err
+		return startTime, headerWritten, ip, proxy, username, err
 	}
-	return startTime, headerWritten, ip, username, nil
+	return startTime, headerWritten, ip, proxy, username, nil
 }
 
 func (e *encoder) handleConnect(storage storage.Writer, msg message.Message, startTime int64, country string, username *string) (string, *string) {
 	payload := msg.Payload.(message.PayloadConnect)
 	ip := payload.RemoteAddr
-	storage.SetMetadata(startTime/1000000000, ip, country, username)
+	var proxy *string
+	if payload.ProxyAddr != "" {
+		proxy = &payload.ProxyAddr
+	}
+	storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	return ip, username
 }
 
-func (e *encoder) handleAuthPasswordSuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, country string) (string, *string) {
+func (e *encoder) handleAuthPasswordSuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, proxy *string, country string) (string, *string) {
 	payload := msg.Payload.(message.PayloadAuthPassword)
 	username := &payload.Username
-	storage.SetMetadata(startTime/1000000000, ip, country, username)
+	storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	return ip, username
 }
 
-func (e *encoder) handleAuthPubkeySuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, country string) (string, *string) {
+func (e *encoder) handleAuthPubkeySuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, proxy *string, country string) (string, *string) {
 	payload := msg.Payload.(message.PayloadAuthPubKey)
 	username := &payload.Username
-	storage.SetMetadata(startTime/1000000000, ip, country, username)
+	storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	return ip, username
 }
 
-func (e *encoder) handleHandshakeSuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, country string) (string, *string) {
+func (e *encoder) handleHandshakeSuccessful(storage storage.Writer, msg message.Message, startTime int64, ip string, proxy *string, country string) (string, *string) {
 	payload := msg.Payload.(message.PayloadHandshakeSuccessful)
 	username := &payload.Username
-	storage.SetMetadata(startTime/1000000000, ip, country, username)
+	storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	return ip, username
 }
 

internal/auditlog/codec/asciinema/encoder_test.go

@@ -19,6 +19,7 @@ type writer struct {
 	data      bytes.Buffer
 	startTime int64
 	sourceIP  string
+	proxyIP   *string
 	username  *string
 	wait      chan bool
 	country   string
@@ -43,9 +44,10 @@ func (w *writer) waitForClose() {
 	<-w.wait
 }
 
-func (w *writer) SetMetadata(startTime int64, sourceIP string, country string, username *string) {
+func (w *writer) SetMetadata(startTime int64, sourceIP string, proxyIP *string, country string, username *string) {
 	w.startTime = startTime
 	w.sourceIP = sourceIP
+	w.proxyIP = proxyIP
 	w.username = username
 	w.country = country
 }

internal/auditlog/codec/binary/encode.go

@@ -48,6 +48,7 @@ func (e *encoder) Encode(messages <-chan message.Message, storage storage.Writer
 
 	startTime := int64(0)
 	var ip = ""
+	var proxy *string
 	var country = "XX"
 	var username *string
 	for {
@@ -58,7 +59,7 @@ func (e *encoder) Encode(messages <-chan message.Message, storage storage.Writer
 		if startTime == 0 {
 			startTime = msg.Timestamp
 		}
-		ip, country, username = e.storeMetadata(msg, storage, startTime, ip, country, username)
+		ip, proxy, country, username = e.storeMetadata(msg, storage, startTime, ip, proxy, country, username)
 		if err := encoder.Encode(&msg); err != nil {
 			return fmt.Errorf("failed to encode audit log message (%w)", err)
 		}
@@ -83,28 +84,32 @@ func (e *encoder) storeMetadata(
 	storage storage.Writer,
 	startTime int64,
 	ip string,
+	proxy *string,
 	country string,
 	username *string,
-) (string, string, *string) {
+) (string, *string, string, *string) {
 	switch msg.MessageType {
 	case message.TypeConnect:
-		remoteAddr := msg.Payload.(message.PayloadConnect).RemoteAddr
-		ip = remoteAddr
+		payload := msg.Payload.(message.PayloadConnect)
+		ip = payload.RemoteAddr
+		if payload.ProxyAddr != "" {
+			proxy = &payload.ProxyAddr
+		}
 		country := e.geoIPProvider.Lookup(net.ParseIP(ip))
-		storage.SetMetadata(startTime/1000000000, ip, country, username)
+		storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	case message.TypeAuthPasswordSuccessful:
 		u := msg.Payload.(message.PayloadAuthPassword).Username
 		username = &u
-		storage.SetMetadata(startTime/1000000000, ip, country, username)
+		storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	case message.TypeAuthPubKeySuccessful:
 		payload := msg.Payload.(message.PayloadAuthPubKey)
 		username = &payload.Username
-		storage.SetMetadata(startTime/1000000000, ip, country, username)
+		storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	case message.TypeHandshakeSuccessful:
 		payload := msg.Payload.(message.PayloadHandshakeSuccessful)
 		username = &payload.Username
-		storage.SetMetadata(startTime/1000000000, ip, country, username)
+		storage.SetMetadata(startTime/1000000000, ip, proxy, country, username)
 	}
 
-	return ip, country, username
+	return ip, proxy, country, username
 }

internal/auditlog/codec/proxy.go

@@ -23,6 +23,6 @@ func (s *storageWriterProxy) Close() error {
 	return s.backend.Close()
 }
 
-func (s *storageWriterProxy) SetMetadata(_ int64, _ string, _ string, _ *string) {
+func (s *storageWriterProxy) SetMetadata(_ int64, _ string, _ *string, _ string, _ *string) {
 	// No metadata storage
 }

internal/auditlog/logger.go

@@ -12,7 +12,7 @@ import (
 type Logger interface {
 	// OnConnect creates an audit log message for a new connection and simultaneously returns a connection object for
 	//           connection-specific messages
-	OnConnect(connectionID message.ConnectionID, ip net.TCPAddr) (Connection, error)
+	OnConnect(connectionID message.ConnectionID, ip net.TCPAddr, proxy *net.TCPAddr) (Connection, error)
 	// Shutdown triggers all failing uploads to cancel, waits for all currently running uploads to finish, then returns.
 	// When the shutdownContext expires it will do its best to immediately upload any running background processes.
 	Shutdown(shutdownContext context.Context)

internal/auditlog/logger_empty.go

@@ -95,7 +95,7 @@ func (e *empty) OnNewChannelSuccess(_ message.ChannelID, _ string) Channel {
 	return e
 }
 
-func (e *empty) OnConnect(_ message.ConnectionID, _ net.TCPAddr) (Connection, error) {
+func (e *empty) OnConnect(_ message.ConnectionID, _ net.TCPAddr, _ *net.TCPAddr) (Connection, error) {
 	return e, nil
 }
 

internal/auditlog/logger_impl.go

@@ -55,7 +55,7 @@ func (l *loggerImplementation) Shutdown(shutdownContext context.Context) {
 
 //region Connection
 
-func (l *loggerImplementation) OnConnect(connectionID message.ConnectionID, ip net.TCPAddr) (Connection, error) {
+func (l *loggerImplementation) OnConnect(connectionID message.ConnectionID, ip net.TCPAddr, proxy *net.TCPAddr) (Connection, error) {
 	name := string(connectionID)
 	writer, err := l.storage.OpenWriter(name)
 	if err != nil {
@@ -76,12 +76,17 @@ func (l *loggerImplementation) OnConnect(connectionID message.ConnectionID, ip n
 			l.logger.Emergency(err)
 		}
 	}()
+	proxyAddr := ""
+	if proxy != nil {
+		proxyAddr = proxy.IP.String()
+	}
 	conn.log(message.Message{
 		ConnectionID: connectionID,
 		Timestamp:    time.Now().UnixNano(),
 		MessageType:  message.TypeConnect,
 		Payload: message.PayloadConnect{
 			RemoteAddr: ip.IP.String(),
+			ProxyAddr:  proxyAddr,
 			Country:    l.geoIPLookup.Lookup(ip.IP),
 		},
 		ChannelID: nil,

internal/auditlog/logger_test.go

@@ -203,6 +203,7 @@ func TestConnect(t *testing.T) {
 			Port: 2222,
 			Zone: "",
 		},
+		nil,
 	)
 	if err != nil {
 		assert.Fail(t, "failed to send connect message to logger", err)
@@ -246,6 +247,7 @@ func TestAuth(t *testing.T) {
 			Port: 2222,
 			Zone: "",
 		},
+		nil,
 	)
 	assert.Nil(t, err)
 	connection.OnAuthPassword("foo", []byte("bar"))

internal/auditlog/storage/file/struct.go

@@ -74,5 +74,5 @@ func (w *writer) Close() error {
 	return w.file.Close()
 }
 
-func (w *writer) SetMetadata(_ int64, _ string, _ string, _ *string) {
+func (w *writer) SetMetadata(_ int64, _ string, _ *string, _ string, _ *string) {
 }