Add SessionPolicyCallback

Closes #7

Author: belak

context.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/hex"
 	"net"
+	"sync"
 
 	gossh "golang.org/x/crypto/ssh"
 )
@@ -59,9 +60,11 @@ var (
 // Context is a package specific context interface. It exposes connection
 // metadata and allows new values to be easily written to it. It's used in
 // authentication handlers and callbacks, and its underlying context.Context is
-// exposed on Session in the session Handler.
+// exposed on Session in the session Handler. This also embeds a sync.Locker so
+// values can safely be set between different sessions.
 type Context interface {
 	context.Context
+	sync.Locker
 
 	// User returns the username used when establishing the SSH connection.
 	User() string
@@ -90,11 +93,12 @@ type Context interface {
 
 type sshContext struct {
 	context.Context
+	sync.Mutex
 }
 
 func newContext(srv *Server) (*sshContext, context.CancelFunc) {
 	innerCtx, cancel := context.WithCancel(context.Background())
-	ctx := &sshContext{innerCtx}
+	ctx := &sshContext{innerCtx, sync.Mutex{}}
 	ctx.SetValue(ContextKeyServer, srv)
 	perms := &Permissions{&gossh.Permissions{}}
 	ctx.SetValue(ContextKeyPermissions, perms)

server.go

@@ -29,6 +29,7 @@ type Server struct {
 	PtyCallback                 PtyCallback                 // callback for allowing PTY sessions, allows all if nil
 	ConnCallback                ConnCallback                // optional callback for wrapping net.Conn before handling
 	LocalPortForwardingCallback LocalPortForwardingCallback // callback for allowing local port forwarding, denies all if nil
+	SessionPolicyCallback       SessionPolicyCallback
 
 	IdleTimeout time.Duration // connection timeout when no activity, none if empty
 	MaxTimeout  time.Duration // absolute connection timeout, none if empty

session.go

@@ -84,30 +84,32 @@ func sessionHandler(srv *Server, conn *gossh.ServerConn, newChan gossh.NewChanne
 		return
 	}
 	sess := &session{
-		Channel: ch,
-		conn:    conn,
-		handler: srv.Handler,
-		ptyCb:   srv.PtyCallback,
-		ctx:     ctx,
+		Channel:   ch,
+		conn:      conn,
+		handler:   srv.Handler,
+		ptyCb:     srv.PtyCallback,
+		sessPolCb: srv.SessionPolicyCallback,
+		ctx:       ctx,
 	}
 	sess.handleRequests(reqs)
 }
 
 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
-	cmd     []string
-	ctx     Context
-	sigCh   chan<- Signal
-	sigBuf  []Signal
+	conn      *gossh.ServerConn
+	handler   Handler
+	handled   bool
+	exited    bool
+	pty       *Pty
+	winch     chan Window
+	env       []string
+	ptyCb     PtyCallback
+	sessPolCb SessionPolicyCallback
+	cmd       []string
+	ctx       Context
+	sigCh     chan<- Signal
+	sigBuf    []Signal
 }
 
 func (sess *session) Write(p []byte) (n int, err error) {
@@ -209,12 +211,22 @@ func (sess *session) handleRequests(reqs <-chan *gossh.Request) {
 				req.Reply(false, nil)
 				continue
 			}
-			sess.handled = true
-			req.Reply(true, nil)
 
 			var payload = struct{ Value string }{}
 			gossh.Unmarshal(req.Payload, &payload)
 			sess.cmd, _ = shlex.Split(payload.Value, true)
+
+			// If there's a session policy callback, we need to confirm before
+			// accepting the session.
+			if sess.sessPolCb != nil && !sess.sessPolCb(sess, req.Type) {
+				sess.cmd = nil
+				req.Reply(false, nil)
+				continue
+			}
+
+			sess.handled = true
+			req.Reply(true, nil)
+
 			go func() {
 				sess.handler(sess)
 				sess.Exit(0)

ssh.go

@@ -42,6 +42,9 @@ type PasswordHandler func(ctx Context, password string) bool
 // PtyCallback is a hook for allowing PTY sessions.
 type PtyCallback func(ctx Context, pty Pty) bool
 
+// SessionPolicyCallback is a callback for allowing SSH sessions.
+type SessionPolicyCallback func(sess Session, requestType string) bool
+
 // ConnCallback is a hook for new connections before handling.
 // It allows wrapping for timeouts and limiting by returning
 // the net.Conn that will be used as the underlying connection.