From 30cc426c915ffa3acd9dbf018efa7ef8f36c7ae8 Mon Sep 17 00:00:00 2001
From: Vasiliy Tolstov <v.tolstov@unistack.org>
Date: Thu, 24 Aug 2023 12:24:03 +0300
Subject: [PATCH] xid fix

Signed-off-by: Vasiliy Tolstov <v.tolstov@unistack.org>
---
 requestid.go | 56 +++++++++++++++++++++++++++++++---------------------
 1 file changed, 34 insertions(+), 22 deletions(-)

diff --git a/requestid.go b/requestid.go
index e45fd44..bdba1a0 100644
--- a/requestid.go
+++ b/requestid.go
@@ -17,32 +17,44 @@ var DefaultMetadataKey = textproto.CanonicalMIMEHeaderKey("x-request-id")
 
 // DefaultMetadataFunc wil be used if user not provide own func to fill metadata
 var DefaultMetadataFunc = func(ctx context.Context) (context.Context, error) {
-	imd, ok := metadata.FromIncomingContext(ctx)
-	if !ok {
-		imd = metadata.New(1)
-	}
-	omd, ok := metadata.FromOutgoingContext(ctx)
-	if !ok {
-		omd = metadata.New(1)
-	}
-	v, iok := imd.Get(DefaultMetadataKey)
-	if iok {
-		if _, ook := omd.Get(DefaultMetadataKey); ook {
-			return ctx, nil
-		}
-	}
-	if !iok {
-		uid, err := id.New()
+	var xid string
+	var err error
+	var ook, iok bool
+
+	if _, ok := ctx.Value(XRequestIDKey).(string); !ok {
+		xid, err = id.New()
 		if err != nil {
 			return ctx, err
 		}
-		v = uid
+		ctx = context.WithValue(ctx, XRequestIDKey, xid)
+	}
+
+	imd, ok := metadata.FromIncomingContext(ctx)
+	if !ok {
+		imd = metadata.New(1)
+		imd.Set(DefaultMetadataKey, xid)
+	} else if _, ok = imd.Get(DefaultMetadataKey); !ok {
+		imd.Set(DefaultMetadataKey, xid)
+	} else {
+		iok = true
+	}
+
+	omd, ok := metadata.FromOutgoingContext(ctx)
+	if !ok {
+		omd = metadata.New(1)
+		omd.Set(DefaultMetadataKey, xid)
+	} else if _, ok = omd.Get(DefaultMetadataKey); !ok {
+		omd.Set(DefaultMetadataKey, xid)
+	} else {
+		ook = true
+	}
+
+	if !iok {
+		ctx = metadata.NewIncomingContext(ctx, imd)
+	}
+	if !ook {
+		ctx = metadata.NewOutgoingContext(ctx, omd)
 	}
-	imd.Set(DefaultMetadataKey, v)
-	omd.Set(DefaultMetadataKey, v)
-	ctx = context.WithValue(ctx, XRequestIDKey, v)
-	ctx = metadata.NewIncomingContext(ctx, imd)
-	ctx = metadata.NewOutgoingContext(ctx, omd)
 
 	return ctx, nil
 }