1 files changed, 41 insertions, 0 deletions
diff --git a/modules/context/private.go b/modules/context/private.go
index dfefa1d2f0..3e31a7e7d8 100644
--- a/modules/context/private.go
+++ b/modules/context/private.go
@@ -6,12 +6,42 @@ package context
 
 import (
 	"context"
+	"fmt"
 	"net/http"
+	"time"
+
+	"code.gitea.io/gitea/modules/graceful"
+	"code.gitea.io/gitea/modules/process"
 )
 
 // PrivateContext represents a context for private routes
 type PrivateContext struct {
 	*Context
+	Override context.Context
+}
+
+// Deadline is part of the interface for context.Context and we pass this to the request context
+func (ctx *PrivateContext) Deadline() (deadline time.Time, ok bool) {
+	if ctx.Override != nil {
+		return ctx.Override.Deadline()
+	}
+	return ctx.Req.Context().Deadline()
+}
+
+// Done is part of the interface for context.Context and we pass this to the request context
+func (ctx *PrivateContext) Done() <-chan struct{} {
+	if ctx.Override != nil {
+		return ctx.Override.Done()
+	}
+	return ctx.Req.Context().Done()
+}
+
+// Err is part of the interface for context.Context and we pass this to the request context
+func (ctx *PrivateContext) Err() error {
+	if ctx.Override != nil {
+		return ctx.Override.Err()
+	}
+	return ctx.Req.Context().Err()
 }
 
 var (
@@ -39,7 +69,18 @@ func PrivateContexter() func(http.Handler) http.Handler {
 				},
 			}
 			ctx.Req = WithPrivateContext(req, ctx)
+			ctx.Data["Context"] = ctx
 			next.ServeHTTP(ctx.Resp, ctx.Req)
 		})
 	}
 }
+
+// OverrideContext overrides the underlying request context for Done() etc.
+// This function should be used when there is a need for work to continue even if the request has been cancelled.
+// Primarily this affects hook/post-receive and hook/proc-receive both of which need to continue working even if
+// the underlying request has timed out from the ssh/http push
+func OverrideContext(ctx *PrivateContext) (cancel context.CancelFunc) {
+	// We now need to override the request context as the base for our work because even if the request is cancelled we have to continue this work
+	ctx.Override, _, cancel = process.GetManager().AddContext(graceful.GetManager().HammerContext(), fmt.Sprintf("PrivateContext: %s", ctx.Req.RequestURI))
+	return
+}