* fix issue where function executions were always crashing under normal circumstances when connection was severed from the server side.

* set the number of function retries from default of 2 to 0 in order to prevent unexpected function re-executions (like for issue above which was resulting in more function executions than expected).
* close connections and don't rely on the crashing of Lambda sandbox to do this.
* pick a dynamic listening port for socks proxy on Lambda side in order to avoid port conflicts.
* add a UUID to Lambda payload for tracing function executions

Author: dan-v

cmd/awslambdaproxy/cmd/run.go

@@ -19,7 +19,7 @@ var (
 	debugProxy 						     bool
 	sshUser, sshPort, regions, listeners string
 	// Max execution time on lambda is 900 seconds currently
-	lambdaMaxFrequency  = time.Duration(890 * time.Second) // leave 10 seconds of leeway
+	lambdaMaxFrequency = time.Duration(time.Second * 860)
 	lambdaMinMemorySize = 128
 	lambdaMaxMemorySize = 1536
 )
@@ -82,7 +82,7 @@ func init() {
 
 	runCmd.Flags().StringVarP(&regions, "regions", "r", "us-west-2", "Regions to "+
 		"run proxy.")
-	runCmd.Flags().DurationVarP(&frequency, "frequency", "f", time.Duration(time.Second*860), "Frequency "+
+	runCmd.Flags().DurationVarP(&frequency, "frequency", "f", lambdaMaxFrequency, "Frequency "+
 		"to execute Lambda function.  Maximum is "+lambdaMaxFrequency.String()+". If multiple "+
 		"lambda-regions are specified, this will cause traffic to rotate round robin at the interval "+
 		"specified here.")

go.mod

@@ -13,6 +13,7 @@ require (
 	github.com/ginuerzh/gost v0.0.0-20200414134316-6e46ac03c7a7
 	github.com/go-ini/ini v1.32.0 // indirect
 	github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b
+	github.com/google/uuid v1.1.1
 	github.com/gorilla/websocket v1.4.2 // indirect
 	github.com/hashicorp/yamux v0.0.0-20190923154419-df201c70410d
 	github.com/inconshreveable/mousetrap v1.0.0 // indirect

go.sum

@@ -103,6 +103,8 @@ github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5a
 github.com/google/gofuzz v1.0.0/go.mod h1:dBl0BpW6vV/+mYPU4Po3pmUjxk6FQPldtuIdl/M65Eg=
 github.com/google/gopacket v1.1.17 h1:rMrlX2ZY2UbvT+sdz3+6J+pp2z+msCq9MxTU6ymxbBY=
 github.com/google/gopacket v1.1.17/go.mod h1:UdDNZ1OO62aGYVnPhxT1U6aI7ukYtA/kB8vaU0diBUM=
+github.com/google/uuid v1.1.1 h1:Gkbcsh/GbpXz7lPftLA3P6TYMwjCLYm83jiFQZF/3gY=
+github.com/google/uuid v1.1.1/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1 h1:EGx4pi6eqNxGaHF6qqu48+N2wcFQ5qg5FXgOdqsJ5d8=
 github.com/gopherjs/gopherjs v0.0.0-20181017120253-0766667cb4d1/go.mod h1:wJfORRmW1u3UXTncJ5qlYoELFm8eSnnEO6hX4iZ3EWY=
 github.com/gorilla/websocket v1.4.0/go.mod h1:E7qHFY5m1UJ88s3WnNqhKjPHQ0heANvMoAMk2YaljkQ=

lambda/lambda.go

@@ -11,6 +11,7 @@ import (
 const privateKeyFile = "/tmp/privatekey"
 
 type Request struct {
+	UUID    string `json:"UUID"`
 	Address string `json:"ConnectBackAddress"`
 	SSHPort string `json:"SSHPort"`
 	SSHKey  string `json:"SSHKey"`
@@ -23,20 +24,23 @@ type Response struct {
 }
 
 func Handler(request Request) (Response, error) {
+	log.Printf("Processing request UUID=%v\n", request.UUID)
 	sshKeyData := []byte(request.SSHKey)
 	err := ioutil.WriteFile(privateKeyFile, sshKeyData, 0600)
 	if err != nil {
 		log.Fatal("Failed to write SSH key to disk. ", err)
 	}
 
-	log.Println("Starting lambdaProxyServer")
+	log.Println("Starting proxy server")
 	lambdaProxyServer := startLambdaProxyServer()
+	defer lambdaProxyServer.close()
 
-	log.Println("Establishing tunnel connection to", request.Address)
+	log.Printf("Establishing ssh tunnel connection to %v\n", request.Address)
 	lambdaTunnelConnection, err := setupLambdaTunnelConnection(request.Address, request.SSHPort, request.SSHUser, privateKeyFile)
 	if err != nil {
-		log.Fatal("Failed to establish connection to "+request.Address, err)
+		log.Fatalf("Failed to establish connection to %v: %v\n", request.Address, err)
 	}
+	defer lambdaTunnelConnection.close()
 
 	log.Println("Starting lambdaDataCopyManager")
 	dataCopyManager := newLambdaDataCopyManager(lambdaProxyServer, lambdaTunnelConnection)

lambda/lambdadatacopy.go

@@ -4,7 +4,6 @@ import (
 	"io"
 	"log"
 	"net"
-	"os"
 	"sync"
 	"time"
 )
@@ -18,18 +17,18 @@ func (l *lambdaDataCopyManager) run() {
 	for {
 		proxySocketConn, proxySocketErr := net.Dial("tcp", l.lambdaProxyServer.port)
 		if proxySocketErr != nil {
-			log.Println("Failed to open connection to proxy", proxySocketErr)
+			log.Printf("Failed to open connection to proxy: %v\n", proxySocketErr)
 			time.Sleep(time.Second)
 			continue
 		}
-		log.Println("Started connection to proxy on port " + l.lambdaProxyServer.port)
+		log.Printf("Opened local connection to proxy on port %v\n", l.lambdaProxyServer.port)
 
 		tunnelStream, tunnelErr := l.lambdaTunnelConnection.sess.Accept()
 		if tunnelErr != nil {
-			log.Println("Failed to start stream inside session", tunnelErr)
-			os.Exit(1)
+			log.Printf("Failed to start new stream: %v. Exiting function.\n", tunnelErr)
+			return
 		}
-		log.Println("Started stream inside session")
+		log.Println("Started new stream")
 
 		go bidirectionalCopy(tunnelStream, proxySocketConn)
 	}

lambda/lambdaproxyserver.go

@@ -1,30 +1,52 @@
 package main
 
 import (
+	"fmt"
 	"log"
+	"net"
 
 	"github.com/ginuerzh/gost"
 )
 
 type lambdaProxyServer struct {
 	port string
+	ln gost.Listener
+	server *gost.Server
 }
 
 func (l *lambdaProxyServer) run() {
-	ln, err := gost.TCPListener(l.port)
+	ln, err := gost.TCPListener(":0")
 	if err != nil {
 		log.Fatal(err)
 	}
+	l.ln = ln
+	l.port = fmt.Sprintf(":%v", ln.Addr().(*net.TCPAddr).Port)
+
 	h := gost.AutoHandler()
 	s := &gost.Server{Listener: ln}
-	log.Fatal(s.Serve(h))
+	l.server = s
+
+	err = s.Serve(h)
+	if err != nil {
+		log.Printf("Server is now exiting: %v\n", err)
+	}
 }
 
-func startLambdaProxyServer() *lambdaProxyServer {
-	port := ":8080"
-	server := &lambdaProxyServer{
-		port: port,
+func (l *lambdaProxyServer) close() {
+	log.Println("Closing down server")
+	err := l.server.Close()
+	if err != nil {
+		log.Printf("closing server error: %v\n", err)
 	}
+	log.Println("Closing down listener")
+	err = l.ln.Close()
+	if err != nil {
+		log.Printf("closing listener error: %v\n", err)
+	}
+}
+
+func startLambdaProxyServer() *lambdaProxyServer {
+	server := &lambdaProxyServer{}
 	go server.run()
 	return server
 }

lambda/lambdatunnelconnection.go

@@ -1,18 +1,16 @@
 package main
 
 import (
+	"io/ioutil"
 	"log"
 	"net"
-	"os"
-
-	"io/ioutil"
 
 	"github.com/hashicorp/yamux"
 	"golang.org/x/crypto/ssh"
 )
 
 const (
-	remoteTunnelPort = "localhost:8081"
+	tunnelPortOnRemoteServer = "localhost:8081"
 )
 
 type lambdaTunnelConnection struct {
@@ -38,28 +36,38 @@ func (l *lambdaTunnelConnection) setup() {
 
 	tunnelConn, err := ssh.Dial("tcp", l.tunnelHost, sshConfig)
 	if err != nil {
-		log.Println("Failed to start SSH tunnel to: "+l.tunnelHost+". Error: ", err)
-		os.Exit(1)
+		log.Fatalf("Failed to start SSH tunnel to %v: %v\n", l.tunnelHost, err)
 	}
-	log.Println("Created SSH tunnel to: " + l.tunnelHost)
+	log.Printf("Setup SSH tunnel to tunnelHost=%v\n", l.tunnelHost)
 
-	localConn, err := tunnelConn.Dial("tcp", remoteTunnelPort)
+	localConn, err := tunnelConn.Dial("tcp", tunnelPortOnRemoteServer)
 	if err != nil {
-		log.Println("Failed to create local tunnel to "+remoteTunnelPort, err)
-		os.Exit(1)
+		log.Fatalf("Failed to create connection to tunnelPortOnRemoteServer=%v: %v\n", tunnelPortOnRemoteServer, err)
 	}
 	l.conn = localConn
-	log.Println("Created local tunnel to " + remoteTunnelPort)
+	log.Printf("Setup connection to tunnelPortOnRemoteServer=%v\n", tunnelPortOnRemoteServer)
 
 	tunnelSession, err := yamux.Server(localConn, nil)
 	if err != nil {
-		log.Println("Failed to start session inside tunnel")
-		os.Exit(1)
+		log.Fatalf("Failed to start session inside tunnel: %v\n", err)
 	}
 	log.Println("Started yamux session inside tunnel")
 	l.sess = tunnelSession
 }
 
+func (l *lambdaTunnelConnection) close() {
+	log.Printf("Closing session")
+	err := l.sess.Close()
+	if err != nil {
+		log.Printf("Error closing session: %v", err)
+	}
+	log.Printf("Closing connection")
+	err = l.conn.Close()
+	if err != nil {
+		log.Printf("Error closing connection: %v", err)
+	}
+}
+
 func setupLambdaTunnelConnection(tunnelHost string, sshPort string, sshUsername string,
 	sshPrivateKeyFile string) (*lambdaTunnelConnection, error) {
 	data, err := ioutil.ReadFile(sshPrivateKeyFile)

lambdaexecution.go

@@ -2,6 +2,7 @@ package awslambdaproxy
 
 import (
 	"encoding/json"
+	"github.com/google/uuid"
 	"log"
 	"time"
 
@@ -20,6 +21,7 @@ type lambdaExecutionManager struct {
 }
 
 type lambdaPayload struct {
+	UUID               string
 	ConnectBackAddress string
 	SSHPort            string
 	SSHKey             string
@@ -29,22 +31,50 @@ type lambdaPayload struct {
 func (l *lambdaExecutionManager) run() {
 	log.Println("Using public IP", l.publicIP)
 	log.Println("Lambda execution frequency", l.frequency)
+	count := 0
+	setInvokeConfig := true
 	for {
+		if count > 0 {
+			setInvokeConfig = false
+		}
 		for region := range l.regions {
-			l.executeFunction(region)
+			l.executeFunction(region, setInvokeConfig)
 			time.Sleep(l.frequency)
 		}
+		count++
 	}
 }
 
-func (l *lambdaExecutionManager) executeFunction(region int) error {
+func (l *lambdaExecutionManager) executeFunction(region int, setInvokeConfig bool) error {
 	log.Println("Executing Lambda function in region", l.regions[region])
 	sess, err := getSessionAWS()
 	if err != nil {
 		return err
 	}
+
 	svc := lambda.New(sess, &aws.Config{Region: aws.String(l.regions[region])})
+
+	if setInvokeConfig {
+		maximumEventAgeInSeconds := int64(1800)
+		maximumRetryAttempts := int64(0)
+		log.Printf("Setting invoke configuration maximumRetryAttempts=%v maximumEventAgeInSeconds=%v\n",
+			maximumRetryAttempts, maximumEventAgeInSeconds)
+		_, err = svc.PutFunctionEventInvokeConfig(&lambda.PutFunctionEventInvokeConfigInput{
+			FunctionName:             aws.String(lambdaFunctionName),
+			MaximumEventAgeInSeconds: aws.Int64(maximumEventAgeInSeconds),
+			MaximumRetryAttempts:     aws.Int64(maximumRetryAttempts),
+		})
+		if err != nil {
+			return err
+		}
+	}
+
+	id, err := uuid.NewUUID()
+	if err !=nil {
+		return err
+	}
 	lambdaPayload := lambdaPayload{
+		UUID:               id.String(),
 		ConnectBackAddress: l.publicIP,
 		SSHPort:            l.sshPort,
 		SSHKey:             l.sshKey,
@@ -56,6 +86,8 @@ func (l *lambdaExecutionManager) executeFunction(region int) error {
 		InvocationType: aws.String(lambda.InvocationTypeEvent),
 		Payload:        payload,
 	}
+
+	log.Printf("Invoking Lambda function with UUID=%v\n", lambdaPayload.UUID)
 	_, err = svc.Invoke(params)
 	if err != nil {
 		return errors.Wrap(err, "Failed to execute Lambda function")
@@ -79,7 +111,7 @@ func newLambdaExecutionManager(publicIP string, regions []string, frequency time
 		for {
 			<-onDemandExecution
 			log.Println("Starting new tunnel as existing tunnel failed")
-			executionManager.executeFunction(0)
+			executionManager.executeFunction(0, false)
 			time.Sleep(time.Second * 5)
 		}
 	}()