gommunicator/internal/ssh/server.go

package ssh

import (
	"fmt"
	"log/slog"
	"net"

	"gommunicator/internal/entity"
	"gommunicator/internal/usecase"

	"golang.org/x/crypto/ssh"
	"golang.org/x/term"
)

type Server struct {
	log       *slog.Logger
	uc        *usecase.Usecase
	sshConfig *ssh.ServerConfig
	sshSigner *ssh.Signer
	socket    *net.Listener
}

func New(log *slog.Logger, uc *usecase.Usecase, privateKey []byte) (*Server, error) {
	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}
	config := ssh.ServerConfig{
		NoClientAuth: true,
	}
	config.AddHostKey(signer)

	server := Server{
		sshConfig: &config,
		sshSigner: &signer,
		log:       log,
		uc:        uc,
	}

	return &server, nil
}

func (s *Server) Start(laddr string) error {
	// Once a ServerConfig has been configured, connections can be accepted.
	socket, err := net.Listen("tcp", laddr)
	if err != nil {
		return err
	}
	s.socket = &socket
	s.log.Info("Listening on", "addr", laddr)
	// generate password
	if err := s.newPassword(); err != nil {
		return err
	}
	// handle connnections
	for {
		conn, err := socket.Accept()
		if err != nil {
			s.log.Info("Failed to accept connection, aborting loop", "err", err)
			return err
		}
		// From a standard TCP connection to an encrypted SSH connection
		sshConn, channels, requests, err := ssh.NewServerConn(conn, s.sshConfig)
		if err != nil {
			s.log.Info("Failed to handshake", "err", err)
			continue
		}
		s.log.Info("New connection", "addr", sshConn.RemoteAddr(), "user", sshConn.User())

		go ssh.DiscardRequests(requests)
		go s.handleChannels(channels, sshConn)
	}
}

func (s *Server) handleChannels(channels <-chan ssh.NewChannel, conn *ssh.ServerConn) {
	for ch := range channels {
		if t := ch.ChannelType(); t != "session" {
			ch.Reject(ssh.UnknownChannelType, fmt.Sprintf("unknown channel type: %s", t))
			continue
		}

		channel, requests, err := ch.Accept()
		if err != nil {
			continue
		}

		go func(in <-chan *ssh.Request) {
			defer channel.Close()
			for req := range in {
				// log.Println("Request: ", req.Type, string(req.Payload))
				ok := false
				switch req.Type {
				case "shell":
					// We don't accept any commands (Payload),
					// only the default shell.
					if len(req.Payload) == 0 {
						ok = true
					}
				case "pty-req":
					// Responding 'ok' here will let the client
					// know we have a pty ready for input
					ok = true
				case "window-change":
					continue // no response
				}
				req.Reply(ok, nil)
			}
		}(requests)
		go s.handleShell(channel, conn.User())
	}
}

func (s *Server) handleShell(channel ssh.Channel, username string) {
	defer func() {
		channel.Close()
		s.log.Info("User disconnected", "user", username)
	}()
	// create terminal
	term := term.NewTerminal(channel, fmt.Sprintf("%s > ", username))
	// check password
	err := s.passwordRequest(term)
	if err != nil {
		return
	}
	// add as client
	s.uc.AddMessageClient(entity.NewTerminalClient(term, username))
	// Recieve user input and send to server
	for {
		// get user input
		line, err := term.ReadLine()
		if err != nil {
			break
		}

		if len(line) > 0 {
			if string(line[0]) == "/" {
				switch line {
				case "/exit":
					return
				case "/help":
					term.Write([]byte(helpMessage))
				default:
					term.Write([]byte(helpMessage))
				}
				continue
			}

			s.log.Debug("new message", "user", username, "messg", line)

			// send message to room
			s.uc.SendMessage(line, username)
		}
	}
}