From a462277fdd1e8c53a67eab1cd1b8cc2658a28181 Mon Sep 17 00:00:00 2001
From: Kaleb Elwert <belak@coded.io>
Date: Wed, 19 Jun 2019 01:19:34 -0700
Subject: [PATCH] Add subsystem support

---
 server.go  | 14 ++++++++++
 session.go | 82 +++++++++++++++++++++++++++++++++++++++++-------------
 2 files changed, 77 insertions(+), 19 deletions(-)

diff --git a/server.go b/server.go
index 359f967..9528b6e 100644
--- a/server.go
+++ b/server.go
@@ -15,6 +15,10 @@ import (
 // and ListenAndServeTLS methods after a call to Shutdown or Close.
 var ErrServerClosed = errors.New("ssh: Server closed")
 
+type SubsystemHandler func(s Session)
+
+var DefaultSubsystemHandlers = map[string]SubsystemHandler{}
+
 type RequestHandler func(ctx Context, srv *Server, req *gossh.Request) (ok bool, payload []byte)
 
 var DefaultRequestHandlers = map[string]RequestHandler{}
@@ -57,6 +61,10 @@ type Server struct {
 	// no handlers are enabled.
 	RequestHandlers map[string]RequestHandler
 
+	// SubsystemHandlers are handlers which are similar to the usual SSH command
+	// handlers, but handle named subsystems.
+	SubsystemHandlers map[string]SubsystemHandler
+
 	listenerWg sync.WaitGroup
 	mu         sync.RWMutex
 	listeners  map[net.Listener]struct{}
@@ -95,6 +103,12 @@ func (srv *Server) ensureHandlers() {
 			srv.ChannelHandlers[k] = v
 		}
 	}
+	if srv.SubsystemHandlers == nil {
+		srv.SubsystemHandlers = map[string]SubsystemHandler{}
+		for k, v := range DefaultSubsystemHandlers {
+			srv.SubsystemHandlers[k] = v
+		}
+	}
 }
 
 func (srv *Server) config(ctx Context) *gossh.ServerConfig {
diff --git a/session.go b/session.go
index 6c77d6c..b2270ce 100644
--- a/session.go
+++ b/session.go
@@ -47,6 +47,9 @@ type Session interface {
 	// RawCommand returns the exact command that was provided by the user.
 	RawCommand() string
 
+	// Subsystem returns the subsystem requested by the user.
+	Subsystem() string
+
 	// PublicKey returns the PublicKey used to authenticate. If a public key was not
 	// used it will return nil.
 	PublicKey() PublicKey
@@ -87,12 +90,13 @@ func DefaultSessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.Ne
 		return
 	}
 	sess := &session{
-		Channel:   ch,
-		conn:      conn,
-		handler:   srv.Handler,
-		ptyCb:     srv.PtyCallback,
-		sessReqCb: srv.SessionRequestCallback,
-		ctx:       ctx,
+		Channel:           ch,
+		conn:              conn,
+		handler:           srv.Handler,
+		ptyCb:             srv.PtyCallback,
+		sessReqCb:         srv.SessionRequestCallback,
+		subsystemHandlers: srv.SubsystemHandlers,
+		ctx:               ctx,
 	}
 	sess.handleRequests(reqs)
 }
@@ -100,19 +104,21 @@ func DefaultSessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.Ne
 type session struct {
 	sync.Mutex
 	gossh.Channel
-	conn      *gossh.ServerConn
-	handler   Handler
-	handled   bool
-	exited    bool
-	pty       *Pty
-	winch     chan Window
-	env       []string
-	ptyCb     PtyCallback
-	sessReqCb SessionRequestCallback
-	rawCmd    string
-	ctx       Context
-	sigCh     chan<- Signal
-	sigBuf    []Signal
+	conn              *gossh.ServerConn
+	handler           Handler
+	subsystemHandlers map[string]SubsystemHandler
+	handled           bool
+	exited            bool
+	pty               *Pty
+	winch             chan Window
+	env               []string
+	ptyCb             PtyCallback
+	sessReqCb         SessionRequestCallback
+	rawCmd            string
+	subsystem         string
+	ctx               Context
+	sigCh             chan<- Signal
+	sigBuf            []Signal
 }
 
 func (sess *session) Write(p []byte) (n int, err error) {
@@ -191,6 +197,10 @@ func (sess *session) Command() []string {
 	return append([]string(nil), cmd...)
 }
 
+func (sess *session) Subsystem() string {
+	return sess.subsystem
+}
+
 func (sess *session) Pty() (Pty, <-chan Window, bool) {
 	if sess.pty != nil {
 		return *sess.pty, sess.winch, true
@@ -239,6 +249,40 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 				sess.handler(sess)
 				sess.Exit(0)
 			}()
+		case "subsystem":
+			if sess.handled {
+				req.Reply(false, nil)
+				continue
+			}
+
+			var payload = struct{ Value string }{}
+			gossh.Unmarshal(req.Payload, &payload)
+			sess.subsystem = payload.Value
+
+			// If there's a session policy callback, we need to confirm before
+			// accepting the session.
+			if sess.sessReqCb != nil && !sess.sessReqCb(sess, req.Type) {
+				sess.rawCmd = ""
+				req.Reply(false, nil)
+				continue
+			}
+
+			handler := sess.subsystemHandlers[payload.Value]
+			if handler == nil {
+				handler = sess.subsystemHandlers["default"]
+			}
+			if handler == nil {
+				req.Reply(false, nil)
+				continue
+			}
+
+			sess.handled = true
+			req.Reply(true, nil)
+
+			go func() {
+				handler(sess)
+				sess.Exit(0)
+			}()
 		case "env":
 			if sess.handled {
 				req.Reply(false, nil)