From 7cad98b285ed3622a67ef642201835669420240d Mon Sep 17 00:00:00 2001 From: YR Chen Date: Thu, 15 Aug 2024 22:32:43 +0800 Subject: [PATCH 1/2] Use `log/slog` for logging --- logging.go | 54 -------------------------------------------------- sshmux.go | 58 +++++++++++++++++++++++++++++++++++------------------- 2 files changed, 38 insertions(+), 74 deletions(-) delete mode 100644 logging.go diff --git a/logging.go b/logging.go deleted file mode 100644 index 1732a54..0000000 --- a/logging.go +++ /dev/null @@ -1,54 +0,0 @@ -package main - -import ( - "encoding/json" - "log" - "net" - "time" -) - -type LogMessage struct { - ConnectTime int64 `json:"connect_time"` - DisconnectTime int64 `json:"disconnect_time"` - ClientIp string `json:"remote_ip"` - HostIp string `json:"host_ip"` - Username string `json:"username"` - Authenticated bool `json:"authenticated"` - ClientType string `json:"client_type"` -} - -type Logger struct { - Channel chan LogMessage -} - -func makeLogger(url string) Logger { - channel := make(chan LogMessage, 256) - go func() { - if url == "" { - for range channel { - } - return - } - conn, err := net.Dial("udp", url) - if err != nil { - log.Printf("Logger Dial failed: %s\n", err) - // Drain the channel to avoid blocking - for range channel { - } - return - } - for logMessage := range channel { - jsonMsg, err := json.Marshal(logMessage) - if err != nil { - continue - } - conn.Write(jsonMsg) - } - }() - return Logger{Channel: channel} -} - -func (l Logger) SendLog(logMessage *LogMessage) { - logMessage.DisconnectTime = time.Now().Unix() - l.Channel <- *logMessage -} diff --git a/sshmux.go b/sshmux.go index 392f97a..9a97b95 100644 --- a/sshmux.go +++ b/sshmux.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "log" + "log/slog" "net" "net/netip" "os" @@ -25,7 +26,7 @@ type Server struct { SSHConfig *ssh.ServerConfig ProxyUpstreams []netip.Prefix Authenticator Authenticator - Logger Logger + LogWriter *net.Conn UsernamePolicy UsernamePolicyConfig PasswordPolicy PasswordPolicyConfig } @@ -54,13 +55,21 @@ func makeServer(config Config) (*Server, error) { } proxyUpstreams = append(proxyUpstreams, network) } + var loggerEndpoint *net.Conn = nil + if config.Logger != "" { + conn, err := net.Dial("udp", config.Logger) + if err != nil { + log.Fatalf("Logger Dial failed: %s\n", err) + } + loggerEndpoint = &conn + } sshmux := &Server{ Address: config.Address, Banner: config.Banner, SSHConfig: sshConfig, ProxyUpstreams: proxyUpstreams, Authenticator: makeAuthenticator(config), - Logger: makeLogger(config.Logger), + LogWriter: loggerEndpoint, UsernamePolicy: UsernamePolicyConfig{ InvalidUsername: config.InvalidUsername, InvalidUsernameMessage: config.InvalidUsernameMessage, @@ -105,22 +114,28 @@ func (s *Server) handler(conn net.Conn) { } defer session.Close() - logMessage := LogMessage{ - ConnectTime: time.Now().Unix(), - ClientIp: conn.RemoteAddr().String(), - Username: "", // should be provided by API server - ClientType: "SSH", - Authenticated: true, + var logger *slog.Logger = nil + if s.LogWriter != nil { + logger = slog.New(slog.NewJSONHandler(*s.LogWriter, nil)) } - defer s.Logger.SendLog(&logMessage) + logger = logger.With( + slog.Int64("connect_time", time.Now().Unix()), + slog.String("remote_ip", conn.RemoteAddr().String()), + slog.String("client_type", "SSH"), + ) + defer logger.Info("SSH proxy session", slog.Int64("disconnect_time", time.Now().Unix())) select { case <-s.ctx.Done(): return default: - if err := s.RunPipeSession(session, &logMessage); err != nil { + attrs, err := s.RunPipeSession(session) + if err != nil { log.Println("runPipeSession:", err) } + for _, attr := range attrs { + logger = logger.With(attr) + } } } @@ -282,6 +297,19 @@ func (s *Server) Handshake(session *ssh.PipeSession) error { } } +func (s *Server) RunPipeSession(session *ssh.PipeSession) ([]slog.Attr, error) { + err := s.Handshake(session) + if err != nil { + return make([]slog.Attr, 0), err + } + attrs := []slog.Attr{ + slog.String("username", session.Downstream.User()), + slog.String("host_ip", session.Upstream.RemoteAddr().String()), + slog.Bool("authenticated", true), + } + return attrs, session.RunPipe() +} + func (s *Server) Start() error { // set up TCP listener listener, err := net.Listen("tcp", s.Address) @@ -333,13 +361,3 @@ func (s *Server) Shutdown() { } s.wg.Wait() } - -func (s *Server) RunPipeSession(session *ssh.PipeSession, logMessage *LogMessage) error { - err := s.Handshake(session) - if err != nil { - return err - } - logMessage.Username = session.Downstream.User() - logMessage.HostIp = session.Upstream.RemoteAddr().String() - return session.RunPipe() -} From c9be0c9a07a61b36d1048a31d7537d5b5f228b22 Mon Sep 17 00:00:00 2001 From: YR Chen Date: Sat, 17 Aug 2024 22:20:34 +0800 Subject: [PATCH 2/2] Update sshmux.go --- sshmux.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sshmux.go b/sshmux.go index 9a97b95..c9baf2f 100644 --- a/sshmux.go +++ b/sshmux.go @@ -300,7 +300,7 @@ func (s *Server) Handshake(session *ssh.PipeSession) error { func (s *Server) RunPipeSession(session *ssh.PipeSession) ([]slog.Attr, error) { err := s.Handshake(session) if err != nil { - return make([]slog.Attr, 0), err + return nil, err } attrs := []slog.Attr{ slog.String("username", session.Downstream.User()),