From e5ed1247eddc9b03a4098799a91f4eb6def70e5e Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Fri, 20 Sep 2024 16:22:51 -0400
Subject: [PATCH] rudimentary implementation of 'zrok share reserved' using
 dual-pathed approach (#751)

---
 agent/agentClient/agentClient.go | 24 ++++++++++++++--
 cmd/zrok/shareReserved.go        | 47 ++++++++++++++++++++++++++++++--
 2 files changed, 65 insertions(+), 6 deletions(-)

diff --git a/agent/agentClient/agentClient.go b/agent/agentClient/agentClient.go
index 8459bc60..110eef0a 100644
--- a/agent/agentClient/agentClient.go
+++ b/agent/agentClient/agentClient.go
@@ -3,18 +3,20 @@ package agentClient
 import (
 	"context"
 	"github.com/openziti/zrok/agent/agentGrpc"
+	"github.com/openziti/zrok/build"
 	"github.com/openziti/zrok/environment/env_core"
-	"github.com/openziti/zrok/tui"
+	"github.com/pkg/errors"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
 	"google.golang.org/grpc/resolver"
 	"net"
+	"strings"
 )
 
 func NewClient(root env_core.Root) (client agentGrpc.AgentClient, conn *grpc.ClientConn, err error) {
 	agentSocket, err := root.AgentSocket()
 	if err != nil {
-		tui.Error("error getting agent socket", err)
+		return nil, nil, err
 	}
 
 	opts := []grpc.DialOption{
@@ -26,8 +28,24 @@ func NewClient(root env_core.Root) (client agentGrpc.AgentClient, conn *grpc.Cli
 	resolver.SetDefaultScheme("passthrough")
 	conn, err = grpc.NewClient(agentSocket, opts...)
 	if err != nil {
-		tui.Error("error connecting to agent socket", err)
+		return nil, nil, err
 	}
 
 	return agentGrpc.NewAgentClient(conn), conn, nil
 }
+
+func IsAgentRunning(root env_core.Root) (bool, error) {
+	client, conn, err := NewClient(root)
+	if err != nil {
+		return false, err
+	}
+	defer func() { _ = conn.Close() }()
+	resp, err := client.Version(context.Background(), &agentGrpc.VersionRequest{})
+	if err != nil {
+		return false, nil
+	}
+	if !strings.HasPrefix(resp.GetV(), build.Series) {
+		return false, errors.Errorf("agent reported version '%v'; we expected version '%v'", resp.GetV(), build.Series)
+	}
+	return true, nil
+}
diff --git a/cmd/zrok/shareReserved.go b/cmd/zrok/shareReserved.go
index 45ff1acd..99824521 100644
--- a/cmd/zrok/shareReserved.go
+++ b/cmd/zrok/shareReserved.go
@@ -1,10 +1,13 @@
 package main
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	tea "github.com/charmbracelet/bubbletea"
 	httptransport "github.com/go-openapi/runtime/client"
+	"github.com/openziti/zrok/agent/agentClient"
+	"github.com/openziti/zrok/agent/agentGrpc"
 	"github.com/openziti/zrok/endpoints"
 	"github.com/openziti/zrok/endpoints/drive"
 	"github.com/openziti/zrok/endpoints/proxy"
@@ -13,6 +16,7 @@ import (
 	"github.com/openziti/zrok/endpoints/udpTunnel"
 	"github.com/openziti/zrok/endpoints/vpn"
 	"github.com/openziti/zrok/environment"
+	"github.com/openziti/zrok/environment/env_core"
 	"github.com/openziti/zrok/rest_client_zrok/metadata"
 	"github.com/openziti/zrok/rest_client_zrok/share"
 	"github.com/openziti/zrok/rest_model_zrok"
@@ -51,9 +55,6 @@ func newShareReservedCommand() *shareReservedCommand {
 }
 
 func (cmd *shareReservedCommand) run(_ *cobra.Command, args []string) {
-	shrToken := args[0]
-	var target string
-
 	root, err := environment.LoadRoot()
 	if err != nil {
 		if !panicInstead {
@@ -66,6 +67,25 @@ func (cmd *shareReservedCommand) run(_ *cobra.Command, args []string) {
 		tui.Error("unable to load environment; did you 'zrok enable'?", nil)
 	}
 
+	if !cmd.agent {
+		agent, err := agentClient.IsAgentRunning(root)
+		if err != nil {
+			tui.Error("error checking if agent is running", err)
+		}
+		if agent {
+			cmd.agentShareReserved(args, root)
+		} else {
+			cmd.shareReserved(args, root)
+		}
+	} else {
+		cmd.shareReserved(args, root)
+	}
+}
+
+func (cmd *shareReservedCommand) shareReserved(args []string, root env_core.Root) {
+	shrToken := args[0]
+	var target string
+
 	zrok, err := root.Client()
 	if err != nil {
 		if !panicInstead {
@@ -390,3 +410,24 @@ func (cmd *shareReservedCommand) run(_ *cobra.Command, args []string) {
 		close(requests)
 	}
 }
+
+func (cmd *shareReservedCommand) agentShareReserved(args []string, root env_core.Root) {
+	logrus.Info("starting")
+
+	client, conn, err := agentClient.NewClient(root)
+	if err != nil {
+		tui.Error("error connecting to agent", err)
+	}
+	defer conn.Close()
+
+	shr, err := client.ShareReserved(context.Background(), &agentGrpc.ShareReservedRequest{
+		Token:            args[0],
+		OverrideEndpoint: cmd.overrideEndpoint,
+		Insecure:         cmd.insecure,
+	})
+	if err != nil {
+		tui.Error("error sharing reserved share", err)
+	}
+
+	fmt.Println(shr)
+}