// SetForwardedHeader adds a header to the current thrift context indicating // that the call has been forwarded by another node in the ringpop ring. // This header is used when a remote call is received to determine if forwarding // checks needs to be applied. By not forwarding already forwarded calls we // prevent unbound forwarding in the ring in case of memebership disagreement. func SetForwardedHeader(ctx thrift.Context) thrift.Context { headers := ctx.Headers() if len(headers) == 0 { return thrift.WithHeaders(ctx, staticForwardHeaders) } headers[forwardedHeaderName] = "true" return thrift.WithHeaders(ctx, headers) }
func thriftCall(clientt thrift.TChanClient, headers map[string]string, token string) (*echo.Pong, map[string]string, error) { client := echo.NewTChanEchoClient(clientt) ctx, cancel := thrift.NewContext(time.Second) ctx = thrift.WithHeaders(ctx, headers) defer cancel() pong, err := client.Echo(ctx, &echo.Ping{Beep: token}) return pong, ctx.ResponseHeaders(), err }
func TestSetForwardedHeader(t *testing.T) { ctx, _ := thrift.NewContext(0 * time.Second) ctx = SetForwardedHeader(ctx) if ctx.Headers()["ringpop-forwarded"] != "true" { t.Errorf("ringpop forwarding header is not set") } ctx, _ = thrift.NewContext(0 * time.Second) ctx = thrift.WithHeaders(ctx, map[string]string{ "keep": "this key", }) ctx = SetForwardedHeader(ctx) if ctx.Headers()["ringpop-forwarded"] != "true" { t.Errorf("ringpop forwarding header is not set if there were headers set already") } if ctx.Headers()["keep"] != "this key" { t.Errorf("ringpop forwarding header removed a header that was already present") } }
func TestHasForwardedHeader(t *testing.T) { ctx, _ := thrift.NewContext(0 * time.Second) if HasForwardedHeader(ctx) { t.Errorf("ringpop claimed that the forwarded header was set before it was set") } ctx = SetForwardedHeader(ctx) if !HasForwardedHeader(ctx) { t.Errorf("ringpop was not able to identify that the forwarded header was set") } ctx, _ = thrift.NewContext(0 * time.Second) ctx = thrift.WithHeaders(ctx, map[string]string{ "keep": "this key", }) if HasForwardedHeader(ctx) { t.Errorf("ringpop claimed that the forwarded header was set before it was set in the case of alread present headers") } ctx = SetForwardedHeader(ctx) if !HasForwardedHeader(ctx) { t.Errorf("ringpop was not able to identify that the forwarded header was set in the case of alread present headers") } }
func createContext() (thrift.Context, func()) { ctx, cancel := thrift.NewContext(time.Second) ctx = thrift.WithHeaders(ctx, map[string]string{"user": curUser}) return ctx, cancel }