From e44c36845129f0a90b1066b6c033b4e424d873b4 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Tue, 14 Jan 2020 15:23:28 -0500
Subject: [PATCH] Required tweaks found when testing xgress_proxy_udp. (#28)

---
 xgress/request.go            |  7 ++++---
 xgress_proxy_udp/listener.go | 21 +++++++++++++++------
 2 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/xgress/request.go b/xgress/request.go
index e9d26cf10..1827b8c7a 100644
--- a/xgress/request.go
+++ b/xgress/request.go
@@ -53,8 +53,9 @@ func (r *Request) ToJSON() ([]byte, error) {
 }
 
 type Response struct {
-	Success bool   `json:"scc"`
-	Message string `json:"msg"`
+	Success   bool   `json:"scc"`
+	Message   string `json:"msg"`
+	SessionId string `json:"sessionId"`
 }
 
 func ResponseFromJSON(payload []byte) (*Response, error) {
@@ -193,7 +194,7 @@ func CreateSession(ctrl CtrlChannel, peer Connection, request *Request, bindHand
 	if err = sessionInfo.SendStartEgress(); err != nil {
 		return &Response{Success: false, Message: err.Error()}
 	}
-	return &Response{Success: true}
+	return &Response{Success: true, SessionId: sessionInfo.SessionId.Token}
 }
 
 func BindService(ctrl CtrlChannel, token string, serviceId string) error {
diff --git a/xgress_proxy_udp/listener.go b/xgress_proxy_udp/listener.go
index 4a4314378..0816704cf 100644
--- a/xgress_proxy_udp/listener.go
+++ b/xgress_proxy_udp/listener.go
@@ -77,6 +77,7 @@ func (l *listener) rx() {
 			sessionId := data.source.String()
 			session, found := l.sessions[sessionId]
 			if !found {
+				logrus.Infof("session not found for [%s]", sessionId)
 				session = &packetSession{
 					listener:             l,
 					readC:                make(chan []byte, 10),
@@ -86,16 +87,20 @@ func (l *listener) rx() {
 				}
 				session.MarkActivity()
 				l.sessions[sessionId] = session
+				l.handleConnect(session)
 
-				go l.handleConnect(session)
-
-				if !session.closed {
-					session.readC <- data.buffer
+				if session.sessionState == sessionStateEstablished {
+					logrus.Infof("created session [%s] => [%s]", sessionId, session.sessionId)
+				} else {
+					logrus.Infof("session creation failed [%s]", sessionId)
 				}
+			}
 
-			} else if session.sessionState == sessionStateEstablished {
+			if session.sessionState == sessionStateEstablished {
 				session.MarkActivity()
 				session.readC <- data.buffer
+			} else {
+				logrus.Warnf("dropping")
 			}
 
 		case event := <-l.eventChan:
@@ -115,7 +120,10 @@ func (l *listener) rx() {
 func (l *listener) handleConnect(session *packetSession) {
 	request := &xgress.Request{ServiceId: l.service}
 	response := xgress.CreateSession(l.ctrl, session, request, l.bindHandler, l.options)
-	if !response.Success{
+	if response.Success {
+		session.sessionState = sessionStateEstablished
+		session.sessionId = response.SessionId
+	} else {
 		logrus.Errorf("error creating session (%s)", response.Message)
 		_ = session.Close()
 	}
@@ -201,6 +209,7 @@ type packetSession struct {
 	listener             *listener
 	readC                chan []byte
 	addr                 net.Addr
+	sessionId            string
 	sessionState         sessionState
 	timeoutIntervalNanos int64
 	timeoutNanos         int64