From 882e7c5c10a869fc3bc6c7307a14197967d09090 Mon Sep 17 00:00:00 2001
From: Homin Lee <ff4500@gmail.com>
Date: Tue, 7 Dec 2021 19:44:10 +0900
Subject: [PATCH] Fix bugs on contact header (#70)

* pkg/session: use session's contact for req & resp

* pkg/ua: fix to work behind sip load balancer

- fix bug on send 200 for BYE to wrong destination port
- update contact header to contain ua's address
---
 pkg/session/session.go | 16 +++++-----------
 pkg/ua/ua.go           | 31 +++++++++++++++++++++----------
 2 files changed, 26 insertions(+), 21 deletions(-)

diff --git a/pkg/session/session.go b/pkg/session/session.go
index 06eac61..2cb88f4 100644
--- a/pkg/session/session.go
+++ b/pkg/session/session.go
@@ -258,7 +258,7 @@ func (s *Session) Reject(statusCode sip.StatusCode, reason string) {
 	request := s.request
 	s.Log().Debugf("Reject: Request => %s, body => %s", request.Short(), request.Body())
 	response := sip.NewResponseFromRequest(request.MessageID(), request, statusCode, reason, "")
-	response.AppendHeader(s.localURI.AsContactHeader())
+	response.AppendHeader(s.contact)
 	tx.Respond(response)
 }
 
@@ -324,7 +324,7 @@ func (s *Session) Accept(statusCode sip.StatusCode) {
 		sip.CopyHeaders("Content-Type", request, response)
 	}
 
-	response.AppendHeader(s.localURI.AsContactHeader())
+	response.AppendHeader(s.contact)
 	response.SetBody(s.answer, true)
 
 	s.response = response
@@ -356,7 +356,8 @@ func (s *Session) Provisional(statusCode sip.StatusCode, reason string) {
 	} else {
 		response = sip.NewResponseFromRequest(request.MessageID(), request, statusCode, reason, "")
 	}
-	response.AppendHeader(s.localURI.AsContactHeader())
+	response.AppendHeader(s.contact)
+
 	s.response = response
 	tx.Respond(response)
 }
@@ -381,20 +382,13 @@ func (s *Session) makeRequest(uaType string, method sip.RequestMethod, msgID sip
 	newRequest.AppendHeader(to)
 	newRequest.SetRecipient(s.request.Recipient())
 	sip.CopyHeaders("Via", inviteRequest, newRequest)
+	newRequest.AppendHeader(s.contact)
 
 	if uaType == "UAC" {
-		if contact, ok := s.request.Contact(); ok {
-			newRequest.AppendHeader(contact)
-		}
-
 		if len(inviteRequest.GetHeaders("Route")) > 0 {
 			sip.CopyHeaders("Route", inviteRequest, newRequest)
 		}
 	} else if uaType == "UAS" {
-		if contact, ok := s.response.Contact(); ok {
-			newRequest.AppendHeader(contact)
-		}
-
 		if len(inviteResponse.GetHeaders("Route")) > 0 {
 			sip.CopyHeaders("Route", inviteResponse, newRequest)
 		}
diff --git a/pkg/ua/ua.go b/pkg/ua/ua.go
index f50d530..87ec9b7 100644
--- a/pkg/ua/ua.go
+++ b/pkg/ua/ua.go
@@ -201,12 +201,12 @@ func (ua *UserAgent) handleBye(request sip.Request, tx sip.ServerTransaction) {
 			if received, ok := viaHop.Params.Get("received"); ok && received.String() != "" {
 				host = received.String()
 			}
-			if rport, ok := viaHop.Params.Get("rport"); ok && rport != nil && rport.String() != "" {
+			if viaHop.Port != nil {
+				port = *viaHop.Port
+			} else if rport, ok := viaHop.Params.Get("rport"); ok && rport != nil && rport.String() != "" {
 				if p, err := strconv.Atoi(rport.String()); err == nil {
 					port = sip.Port(uint16(p))
 				}
-			} else if request.Recipient().Port() != nil {
-				port = *request.Recipient().Port()
 			} else {
 				port = sip.DefaultPort(request.Transport())
 			}
@@ -271,8 +271,11 @@ func (ua *UserAgent) handleInvite(request sip.Request, tx sip.ServerTransaction)
 			is.SetState(session.ReInviteReceived)
 			ua.handleInviteState(is, &request, nil, session.ReInviteReceived, &transaction)
 		} else {
-			contact, _ := request.Contact()
-			is := session.NewInviteSession(ua.RequestWithContext, "UAS", contact, request, *callID, transaction, session.Incoming, ua.Log())
+			contactHdr, _ := request.Contact()
+			contactAddr := ua.updateContact2UAAddr(request.Transport(), contactHdr.Address)
+			contactHdr.Address = contactAddr
+
+			is := session.NewInviteSession(ua.RequestWithContext, "UAS", contactHdr, request, *callID, transaction, session.Incoming, ua.Log())
 			ua.iss.Store(*callID, is)
 			is.SetState(session.InviteReceived)
 			ua.handleInviteState(is, &request, nil, session.InviteReceived, &transaction)
@@ -322,12 +325,12 @@ func (ua *UserAgent) RequestWithContext(ctx context.Context, request sip.Request
 	var cts sip.Transaction = tx.(sip.Transaction)
 
 	if request.IsInvite() {
-		callID, ok := request.CallID()
-		if ok {
-
+		if callID, ok := request.CallID(); ok {
 			if _, found := ua.iss.Load(*callID); !found {
-				contact, _ := request.Contact()
-				is := session.NewInviteSession(ua.RequestWithContext, "UAC", contact, request, *callID, cts, session.Outgoing, ua.Log())
+				contactHdr, _ := request.Contact()
+				contactAddr := ua.updateContact2UAAddr(request.Transport(), contactHdr.Address)
+				contactHdr.Address = contactAddr
+				is := session.NewInviteSession(ua.RequestWithContext, "UAC", contactHdr, request, *callID, cts, session.Outgoing, ua.Log())
 				ua.iss.Store(*callID, is)
 				is.ProvideOffer(request.Body())
 				is.SetState(session.InviteSent)
@@ -516,3 +519,11 @@ func (ua *UserAgent) RequestWithContext(ctx context.Context, request sip.Request
 func (ua *UserAgent) Shutdown() {
 	ua.config.SipStack.Shutdown()
 }
+
+func (ua *UserAgent) updateContact2UAAddr(transport string, from sip.ContactUri) sip.ContactUri {
+	stackAddr := ua.config.SipStack.GetNetworkInfo(transport)
+	ret := from.Clone()
+	ret.SetHost(stackAddr.Host)
+	ret.SetPort(stackAddr.Port)
+	return ret
+}