From c47233b8207f64988ac63cdab18e73ade5ac0dea Mon Sep 17 00:00:00 2001
From: Vladimir Vershinin <ghettovoice@gmail.com>
Date: Tue, 27 Aug 2024 11:31:22 +0300
Subject: [PATCH] Add methods to set up callbacks on 3xx ACK and CANCEL request
 building

`OnAck(fn func(sip.Request))` - sets callback on 3xx ACK request building during INVITE transaction;
`OnCancel(fn func(sip.Request))` - sets callback on CANCEL request building during INVITE transaction.
---
 options.go               | 18 ++++++++++++++++--
 server.go                | 13 ++++++++++---
 sip/transaction.go       |  3 +++
 transaction/client_tx.go | 25 +++++++++++++++++++++++++
 4 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/options.go b/options.go
index e79e641..f2b2591 100644
--- a/options.go
+++ b/options.go
@@ -7,8 +7,9 @@ type RequestWithContextOption interface {
 }
 
 type RequestWithContextOptions struct {
-	ResponseHandler func(res sip.Response, request sip.Request)
-	Authorizer      sip.Authorizer
+	ResponseHandler     func(res sip.Response, request sip.Request)
+	Authorizer          sip.Authorizer
+	OnAckFn, OnCancelFn func(sip.Request)
 }
 
 type withResponseHandler struct {
@@ -34,3 +35,16 @@ func (o withAuthorizer) ApplyRequestWithContext(options *RequestWithContextOptio
 func WithAuthorizer(authorizer sip.Authorizer) RequestWithContextOption {
 	return withAuthorizer{authorizer}
 }
+
+type withCallbacks struct {
+	onAckFn, onCancFn func(sip.Request)
+}
+
+func WithCallbacks(onAckFn, onCancFn func(sip.Request)) RequestWithContextOption {
+	return withCallbacks{onAckFn, onCancFn}
+}
+
+func (o withCallbacks) ApplyRequestWithContext(options *RequestWithContextOptions) {
+	options.OnAckFn = o.onAckFn
+	options.OnCancelFn = o.onCancFn
+}
diff --git a/server.go b/server.go
index 771052b..ffafd56 100644
--- a/server.go
+++ b/server.go
@@ -293,14 +293,21 @@ func (srv *server) requestWithContext(
 	attempt int,
 	options ...RequestWithContextOption,
 ) (sip.Response, error) {
+	optionsHash := &RequestWithContextOptions{}
+	for _, opt := range options {
+		opt.ApplyRequestWithContext(optionsHash)
+	}
+
 	tx, err := srv.Request(request)
 	if err != nil {
 		return nil, err
 	}
 
-	optionsHash := &RequestWithContextOptions{}
-	for _, opt := range options {
-		opt.ApplyRequestWithContext(optionsHash)
+	if optionsHash.OnAckFn != nil {
+		tx.OnAck(optionsHash.OnAckFn)
+	}
+	if optionsHash.OnCancelFn != nil {
+		tx.OnCancel(optionsHash.OnCancelFn)
 	}
 
 	txResponses := tx.Responses()
diff --git a/sip/transaction.go b/sip/transaction.go
index ff010d8..6ca3111 100644
--- a/sip/transaction.go
+++ b/sip/transaction.go
@@ -25,4 +25,7 @@ type ClientTransaction interface {
 	Transaction
 	Responses() <-chan Response
 	Cancel() error
+
+	OnAck(fn func(Request))
+	OnCancel(fn func(Request))
 }
diff --git a/transaction/client_tx.go b/transaction/client_tx.go
index f09c13a..c210c43 100644
--- a/transaction/client_tx.go
+++ b/transaction/client_tx.go
@@ -16,6 +16,9 @@ type ClientTx interface {
 	Tx
 	Responses() <-chan sip.Response
 	Cancel() error
+
+	OnAck(fn func(sip.Request))
+	OnCancel(fn func(sip.Request))
 }
 
 type clientTx struct {
@@ -31,6 +34,8 @@ type clientTx struct {
 
 	mu        sync.RWMutex
 	closeOnce sync.Once
+
+	onAckFn, onCancFn func(sip.Request)
 }
 
 func NewClientTx(origin sip.Request, tpl sip.Transport, logger log.Logger) (ClientTx, error) {
@@ -231,11 +236,15 @@ func (tx *clientTx) cancel() {
 
 	tx.mu.RLock()
 	lastResp := tx.lastResp
+	onCanc := tx.onCancFn
 	tx.mu.RUnlock()
 
 	cancelRequest := sip.NewCancelRequest("", tx.Origin(), log.Fields{
 		"sent_at": time.Now(),
 	})
+	if onCanc != nil {
+		onCanc(cancelRequest)
+	}
 	if err := tx.tpl.Send(cancelRequest); err != nil {
 		var lastRespStr string
 		if lastResp != nil {
@@ -264,11 +273,15 @@ func (tx *clientTx) cancel() {
 func (tx *clientTx) ack() {
 	tx.mu.RLock()
 	lastResp := tx.lastResp
+	onAck := tx.onAckFn
 	tx.mu.RUnlock()
 
 	ack := sip.NewAckRequest("", tx.Origin(), lastResp, "", log.Fields{
 		"sent_at": time.Now(),
 	})
+	if onAck != nil {
+		onAck(ack)
+	}
 	err := tx.tpl.Send(ack)
 	if err != nil {
 		tx.Log().WithFields(log.Fields{
@@ -626,6 +639,18 @@ func (tx *clientTx) delete() {
 	tx.mu.Unlock()
 }
 
+func (tx *clientTx) OnAck(fn func(sip.Request)) {
+	tx.mu.Lock()
+	tx.onAckFn = fn
+	tx.mu.Unlock()
+}
+
+func (tx *clientTx) OnCancel(fn func(sip.Request)) {
+	tx.mu.Lock()
+	tx.onCancFn = fn
+	tx.mu.Unlock()
+}
+
 // Define actions
 func (tx *clientTx) act_invite_resend() fsm.Input {
 	tx.Log().Debug("act_invite_resend")