diff --git a/brim/brim.go b/brim/brim.go index 676cbca..96b395a 100644 --- 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 index a3b7134..b6dda7e 100644 --- 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 index 0f00d4d..f9956f1 100644 --- 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