From 6c9a651a0871be1d64556db45ede916882fa5913 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Tue, 18 Apr 2023 11:58:07 -0400
Subject: [PATCH] tunnelFrontend elaboration (#170)

---
 cmd/zrok/accessPrivateTunnel.go      | 35 +++++++++++++++++++++
 endpoints/tunnelFrontend/frontend.go | 46 ++++++++++++++++++++++++----
 2 files changed, 75 insertions(+), 6 deletions(-)

diff --git a/cmd/zrok/accessPrivateTunnel.go b/cmd/zrok/accessPrivateTunnel.go
index 69f138c0..08db5e1c 100644
--- a/cmd/zrok/accessPrivateTunnel.go
+++ b/cmd/zrok/accessPrivateTunnel.go
@@ -1,7 +1,13 @@
 package main
 
 import (
+	httptransport "github.com/go-openapi/runtime/client"
 	"github.com/openziti/zrok/endpoints/tunnelFrontend"
+	"github.com/openziti/zrok/rest_client_zrok/share"
+	"github.com/openziti/zrok/rest_model_zrok"
+	"github.com/openziti/zrok/tui"
+	"github.com/openziti/zrok/zrokdir"
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"time"
 )
@@ -28,6 +34,35 @@ func newAccessPrivateTunnelCommand() *accessPrivateTunnelCommand {
 }
 
 func (cmd *accessPrivateTunnelCommand) run(_ *cobra.Command, args []string) {
+	zrd, err := zrokdir.Load()
+	if err != nil {
+		tui.Error("unable to load zrokdir", err)
+	}
+
+	if zrd.Env == nil {
+		tui.Error("unable to load environment; did you 'zrok enable'?", nil)
+	}
+
+	zrok, err := zrd.Client()
+	if err != nil {
+		tui.Error("unable to create zrok client", err)
+	}
+
+	auth := httptransport.APIKeyAuth("X-TOKEN", "header", zrd.Env.Token)
+	req := share.NewAccessParams()
+	req.Body = &rest_model_zrok.AccessRequest{
+		ShrToken: args[0],
+		EnvZID:   zrd.Env.ZId,
+	}
+	accessResp, err := zrok.Share.Access(req, auth)
+	if err != nil {
+		if !panicInstead {
+			tui.Error("unable to access", err)
+		}
+		panic(err)
+	}
+	logrus.Infof("allocated frontend '%v'", accessResp.Payload.FrontendToken)
+
 	fe, err := tunnelFrontend.NewFrontend(&tunnelFrontend.Config{
 		BindAddress:  cmd.bindAddress,
 		IdentityName: "backend",
diff --git a/endpoints/tunnelFrontend/frontend.go b/endpoints/tunnelFrontend/frontend.go
index 467ef704..2ad5f8ad 100644
--- a/endpoints/tunnelFrontend/frontend.go
+++ b/endpoints/tunnelFrontend/frontend.go
@@ -3,6 +3,7 @@ package tunnelFrontend
 import (
 	"github.com/openziti/sdk-golang/ziti"
 	"github.com/openziti/sdk-golang/ziti/config"
+	"github.com/openziti/sdk-golang/ziti/edge"
 	"github.com/openziti/transport/v2"
 	"github.com/openziti/zrok/model"
 	"github.com/openziti/zrok/zrokdir"
@@ -62,16 +63,49 @@ func (f *Frontend) Stop() {
 }
 
 func (f *Frontend) accept(conn transport.Conn) {
+	if zConn, err := f.zCtx.Dial(f.cfg.ShrToken); err == nil {
+		go f.rxer(conn, zConn)
+		go f.txer(conn, zConn)
+		logrus.Infof("accepted '%v' <=> '%v'", conn.RemoteAddr(), zConn.RemoteAddr())
+	} else {
+		logrus.Errorf("error dialing '%v': %v", f.cfg.ShrToken, err)
+		_ = conn.Close()
+	}
+}
+
+func (f *Frontend) rxer(conn transport.Conn, zConn edge.Conn) {
 	buf := make([]byte, 10240)
 	for {
-		n, err := conn.Read(buf)
-		if err != nil {
-			logrus.Errorf("error reading: %v", err)
+		if rxsz, err := conn.Read(buf); err == nil {
+			if txsz, err := zConn.Write(buf[:rxsz]); err == nil {
+				if txsz != rxsz {
+					logrus.Errorf("short write '%v' (%d != %d)", zConn.RemoteAddr(), txsz, rxsz)
+				}
+			} else {
+				logrus.Errorf("error writing '%v': %v", zConn.RemoteAddr(), err)
+				return
+			}
+		} else {
+			logrus.Errorf("read error '%v': %v", zConn.RemoteAddr(), err)
 			return
 		}
-		n, err = conn.Write(buf[:n])
-		if err != nil {
-			logrus.Errorf("error writing: %v", err)
+	}
+}
+
+func (f *Frontend) txer(conn transport.Conn, zConn edge.Conn) {
+	buf := make([]byte, 10240)
+	for {
+		if rxsz, err := zConn.Read(buf); err == nil {
+			if txsz, err := conn.Write(buf[:rxsz]); err == nil {
+				if txsz != rxsz {
+					logrus.Errorf("short write '%v' (%d != %d)'", conn.RemoteAddr(), txsz, rxsz)
+				}
+			} else {
+				logrus.Errorf("error writing '%v': %v", conn.RemoteAddr(), err)
+				return
+			}
+		} else {
+			logrus.Errorf("read error '%v': %v", conn.RemoteAddr(), err)
 			return
 		}
 	}