From f50cd4622a4143d2265c3e548ad2d5346c5bd19a Mon Sep 17 00:00:00 2001
From: jspc <james+git@zero-internet.org.uk>
Date: Wed, 13 Mar 2024 21:32:40 +0000
Subject: [PATCH] Allow for finalisers

It is handy to run functions after a successful run, for things like clean up tasks, or indexing, or whatever
---
 ssh.go | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/ssh.go b/ssh.go
index 87aa0e0..e50ed08 100644
--- a/ssh.go
+++ b/ssh.go
@@ -35,6 +35,7 @@ type PublicKey struct {
 
 type PublicKeyContextKey struct{}
 type UserContextKey struct{}
+type RepositoryContextKey struct{}
 
 const (
 	keyID   = "key-id"
@@ -51,6 +52,7 @@ type SSH struct {
 	PublicKeyLookupFunc    func(ctx context.Context, publicKeyPayload string) (*PublicKey, error)
 	PreLoginFunc           func(ctx context.Context, metadata ssh.ConnMetadata) error
 	AuthoriseOperationFunc func(ctx context.Context, cmd *GitCommand) error
+	FinaliseFunc           func(ctx context.Context)
 }
 
 func NewSSH(config Config) *SSH {
@@ -233,6 +235,9 @@ func (s SSH) handleExecRequest(ctx context.Context, ch ssh.Channel, req *ssh.Req
 		return fmt.Errorf("ssh: command failed: %w", err)
 	}
 
+	ctx = context.WithValue(ctx, RepositoryContextKey{}, gitcmd)
+	s.FinaliseFunc(ctx)
+
 	_, err = ch.SendRequest("exit-status", true, []byte{0, 0, 0, 0})
 
 	return
@@ -283,6 +288,10 @@ func (s SSH) defaultPreLoginFunc(ctx context.Context, metadata ssh.ConnMetadata)
 	return nil
 }
 
+func (s SSH) defaultFinaliseFunc(ctx context.Context) {
+	return
+}
+
 func (s *SSH) setup() error {
 	if s.sshconfig != nil {
 		return nil
@@ -307,6 +316,10 @@ func (s *SSH) setup() error {
 			s.PreLoginFunc = s.defaultPreLoginFunc
 		}
 
+		if s.FinaliseFunc == nil {
+			s.FinaliseFunc = s.defaultFinaliseFunc
+		}
+
 		config.PublicKeyCallback = func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
 			ctx := context.WithValue(context.Background(), UserContextKey{}, conn.User())
 			err := s.PreLoginFunc(ctx, conn)