rebound/ssh/session_handler.go

package ssh

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"strings"
	"text/template"

	_ "embed"

	"github.com/gliderlabs/ssh"
	"github.com/pkg/errors"
)

//go:embed diagram.txt
var asciiDiagram string
var asciiDiagramTmpl = template.Must(template.New("").Parse(asciiDiagram))

func (s *Server) handleSession(sess ssh.Session) {
	ctx := sess.Context()
	sessionID := SessionID(ctx.SessionID())

	s.log("(%s): session opened", sessionID)

	message := `
	Welcome on Rebound !
`

	if _, err := sess.Write([]byte(message)); err != nil {
		s.log("[ERROR] %+v", errors.WithStack(err))
	}

	go s.readClientInput(sess)

	data := s.sessionManager.Get(sessionID, SessionData{
		Type: TypeServiceUnknown,
	})
	s.handleSessionData(sess, data)

	onUpdate, close := s.sessionManager.OnUpdate(sessionID)
	defer close()

	for {
		data, opened := <-onUpdate
		if !opened {
			return
		}

		s.log("(%s): session data updated: %v", sessionID, data)

		if err := s.handleSessionData(sess, data); err != nil {
			if errors.Is(err, io.EOF) {
				return
			}

			s.log("[ERROR] %+v", errors.WithStack(err))

			return
		}
	}
}

func (s *Server) handleSessionData(sess ssh.Session, data SessionData) error {
	switch data.Type {
	case TypeServiceConsumer:
		if err := s.writeConsumerMessage(sess, data); err != nil {
			return errors.WithStack(err)
		}

	case TypeServiceProvider:
		if err := s.writeProviderMessage(sess, data); err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

func (s *Server) writeConsumerMessage(sess ssh.Session, data SessionData) error {
	message := `
	Type Ctrl+C or Ctrl+D to exit.
	`

	if _, err := sess.Write([]byte(message)); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

func (s *Server) writeProviderMessage(sess ssh.Session, data SessionData) error {
	var diagramBuff bytes.Buffer

	hostname := s.opts.PublicHost
	if len(hostname) < 24 {
		halfPadding := (24 - len(hostname)) / 2
		leftPadding := strings.Repeat(" ", halfPadding)
		rightPadding := strings.Repeat(" ", halfPadding)
		hostname = fmt.Sprintf("%s%s%s", leftPadding, hostname, rightPadding)
		if len(hostname) > 24 {
			hostname = hostname[0:23]
		}
	} else if len(hostname) >= 24 {
		hostname = hostname[0:20] + "..."
	}

	log.Printf("'%s'", hostname)

	tmplData := struct {
		Pp       string
		Rp       string
		Hostname string
	}{
		Pp:       fmt.Sprintf("%04d", s.opts.PublicPort),
		Rp:       "<port>",
		Hostname: hostname,
	}

	if err := asciiDiagramTmpl.Execute(&diagramBuff, tmplData); err != nil {
		return errors.WithStack(err)
	}

	message := fmt.Sprintf(`
	You can connect to your tunnel by running in an other terminal:

	ssh -L <port>:127.0.0.1:1 %s@%s -p %d

%s

	Type Ctrl+C or Ctrl+D to exit.
`, data.Token, s.opts.PublicHost, s.opts.PublicPort, diagramBuff.String())

	if _, err := sess.Write([]byte(message)); err != nil {
		return errors.WithStack(err)
	}

	return nil
}

const (
	CtrlC = 3
	CtrlD = 4
)

func (s *Server) readClientInput(sess ssh.Session) {
	sessionID := SessionID(sess.Context().SessionID())
	defer func() {
		s.sessionManager.Remove(sessionID)
	}()

	buff := make([]byte, 1)

	for {
		_, err := sess.Read(buff)
		if err != nil {
			if !errors.Is(err, io.EOF) {
				s.log("[ERROR] %+v", errors.WithStack(err))
			}

			return
		}

		switch buff[0] {
		case CtrlC:
			fallthrough
		case CtrlD:
			sess.Exit(0)
			return

		default:
			s.log("(%s) user input: %v", sessionID, buff)
		}

	}
}