Merge remote-tracking branch 'upstream/main'

trimble-oss · Aug 29, 2024 · 7c50bdd · 7c50bdd
2 parents ee53d50 + 7dc5cfc
commit 7c50bdd
Show file tree

Hide file tree

Showing 3 changed files with 114 additions and 23 deletions.
diff --git a/brim/brim.go b/brim/brim.go
@@ -11,7 +11,9 @@ import (
 	"github.com/trimble-oss/tierceron-hat/cap"
 	"github.com/trimble-oss/tierceron-hat/cap/tap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
 )
 
 func init() {
@@ -33,26 +35,59 @@ func penseQuery(pense string) {
 	penseArray := sha256.Sum256([]byte(penseCode))
 	penseSum := hex.EncodeToString(penseArray[:])
 
-	eyeMap, err := tap.TapWriter(penseSum)
-	if err != nil {
-		log.Fatalf("Failure to communicate: %v", err)
-	}
-	log.Printf("%v", eyeMap)
-
 	conn, err := grpc.Dial("127.0.0.1:1534", grpc.WithTransportCredentials(insecure.NewCredentials()))
 	if err != nil {
 		log.Fatalf("did not connect: %v", err)
 	}
 	defer conn.Close()
+
 	c := cap.NewCapClient(conn)
 
 	// Contact the server and print out its response.
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
 	defer cancel()
 
-	r, err := c.Pense(ctx, &cap.PenseRequest{Pense: penseCode, PenseIndex: pense})
+	var r *cap.PenseReply
+	retry := 0
+
+	for {
+		_, err := c.Pense(ctx, &cap.PenseRequest{Pense: "", PenseIndex: ""})
+		if err != nil {
+			st, ok := status.FromError(err)
+
+			if ok && (retry < 5) && st.Code() == codes.Unavailable {
+				retry = retry + 1
+				continue
+			} else {
+				log.Fatalf("did not connect: %v", err)
+			}
+		} else {
+			break
+		}
+	}
+
+	eyeMap, err := tap.TapWriter(penseSum)
 	if err != nil {
-		log.Fatalf("did not connect: %v", err)
+		log.Fatalf("Failure to communicate: %v", err)
+	}
+	log.Printf("%v", eyeMap)
+
+	retry = 0
+
+	for {
+		r, err = c.Pense(ctx, &cap.PenseRequest{Pense: penseCode, PenseIndex: pense})
+		if err != nil {
+			st, ok := status.FromError(err)
+
+			if ok && (retry < 5) && st.Code() == codes.Unavailable {
+				retry = retry + 1
+				continue
+			} else {
+				log.Fatalf("did not connect: %v", err)
+			}
+		} else {
+			break
+		}
 	}
 
 	log.Println(pense, r.GetPense())

diff --git a/cap/cap.go b/cap/cap.go
@@ -82,6 +82,15 @@ var penseFeatherMemoryMap map[string]string = map[string]string{}
 var penseFeatherPluckMap = cmap.New[bool]()
 var penseFeatherCtlCodeMap = cmap.New[string]()
 
+// CodeSaltGuardFn is expected to return a hex.EncodeToString encoded salt
+type CodeSaltGuardFn func() string
+
+var codeSaltGuardFn CodeSaltGuardFn = nil
+
+func TapInitCodeSaltGuard(csgFn CodeSaltGuardFn) {
+	codeSaltGuardFn = csgFn
+}
+
 func TapServer(address string, opt ...grpc.ServerOption) {
 	lis, err := net.Listen("tcp", address)
 	if err != nil {
@@ -199,7 +208,20 @@ func handleMessage(handshakeCode string, conn *kcp.UDPSession, acceptRemote func
 						message := string(messageBytes)
 						messageParts := strings.Split(message, string(PROTOCOL_DELIM))
 						if messageParts[0] == handshakeCode {
-							if len(messageParts[1]) == 64 {
+							featherCode := messageParts[1]
+							if codeSaltGuardFn != nil {
+								codeSalt := codeSaltGuardFn()
+								if len(codeSalt) > 0 {
+									if strings.HasSuffix(featherCode, codeSalt) {
+										featherCode = strings.TrimSuffix(featherCode, codeSalt)
+									} else {
+										// Invalid
+										featherCode = ""
+									}
+								}
+
+							}
+							if len(featherCode) == 64 {
 								penseFeatherCodeMap.Set(messageParts[1], "")
 							}
 						}

diff --git a/captip/captiplib/captiplib.go b/captip/captiplib/captiplib.go
@@ -15,7 +15,9 @@ import (
 
 	"github.com/trimble-oss/tierceron-hat/cap"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/status"
 )
 
 func init() {
@@ -187,7 +189,37 @@ func randomString(n int) string {
 }
 
 func FeatherQueryCache(featherCtx *cap.FeatherContext, pense string) (string, error) {
-	penseCode := randomString(7 + rand.Intn(7))
+	conn, err := grpc.Dial(*featherCtx.LocalHostAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		return "", err
+	}
+	defer conn.Close()
+
+	c := cap.NewCapClient(conn)
+	// Contact the server and print out its response.
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+
+	var r *cap.PenseReply
+	retry := 0
+
+	for {
+		_, err := c.Pense(ctx, &cap.PenseRequest{Pense: "", PenseIndex: ""})
+		if err != nil {
+			st, ok := status.FromError(err)
+
+			if ok && (retry < 5) && st.Code() == codes.Unavailable {
+				retry = retry + 1
+				continue
+			} else {
+				return "", err
+			}
+		} else {
+			break
+		}
+	}
+
+	penseCode := randomString(12 + rand.Intn(7))
 	penseArray := sha256.Sum256([]byte(penseCode))
 	penseSum := hex.EncodeToString(penseArray[:])
 
@@ -196,20 +228,22 @@ func FeatherQueryCache(featherCtx *cap.FeatherContext, pense string) (string, er
 		return "", featherErr
 	}
 
-	conn, err := grpc.Dial(*featherCtx.LocalHostAddr, grpc.WithTransportCredentials(insecure.NewCredentials()))
-	if err != nil {
-		return "", err
-	}
-	defer conn.Close()
-	c := cap.NewCapClient(conn)
+	retry = 0
 
-	// Contact the server and print out its response.
-	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
-	defer cancel()
+	for {
+		r, err = c.Pense(ctx, &cap.PenseRequest{Pense: penseCode, PenseIndex: pense})
+		if err != nil {
+			st, ok := status.FromError(err)
 
-	r, err := c.Pense(ctx, &cap.PenseRequest{Pense: penseCode, PenseIndex: pense})
-	if err != nil {
-		return "", err
+			if ok && (retry < 5) && st.Code() == codes.Unavailable {
+				retry = retry + 1
+				continue
+			} else {
+				return "", err
+			}
+		} else {
+			break
+		}
 	}
 
 	return r.GetPense(), nil