// 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
}