// 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 (thriftTestHandler) TestException(ctx thrift.Context, arg string) error {
ctx.SetResponseHeaders(ctx.Headers())
switch arg {
case "Xception":
code := int32(1001)
return &gauntlet_tchannel.Xception{ErrorCode: &code, Message: &arg}
case "TException":
// unexpected exception.
return errors.New("great sadness")
default:
return nil
}
}
// Set sets the value for a given key.
func (h *kvHandler) Set(ctx thrift.Context, key, value string) error {
if err := isValidKey(key); err != nil {
return err
}
h.mut.Lock()
defer h.mut.Unlock()
h.vals[key] = value
// Example of how to use response headers. Normally, these values should be passed via result structs.
ctx.SetResponseHeaders(map[string]string{"count": fmt.Sprint(len(h.vals))})
return nil
}
func (thriftTestHandler) TestMultiException(ctx thrift.Context, arg0 string, arg1 string) (*gauntlet_tchannel.Xtruct, error) {
ctx.SetResponseHeaders(ctx.Headers())
structThing := &gauntlet_tchannel.Xtruct{StringThing: &arg1}
switch arg0 {
case "Xception":
code := int32(1001)
message := "This is an Xception"
return nil, &gauntlet_tchannel.Xception{ErrorCode: &code, Message: &message}
case "Xception2":
code := int32(2002)
return nil, &gauntlet_tchannel.Xception2{ErrorCode: &code, StructThing: structThing}
default:
return structThing, nil
}
}
func (thriftTestHandler) TestInsanity(ctx thrift.Context, argument *gauntlet_tchannel.Insanity) (
map[gauntlet_tchannel.UserId]map[gauntlet_tchannel.Numberz]*gauntlet_tchannel.Insanity, error) {
ctx.SetResponseHeaders(ctx.Headers())
result := map[gauntlet_tchannel.UserId]map[gauntlet_tchannel.Numberz]*gauntlet_tchannel.Insanity{
1: {
gauntlet_tchannel.Numberz_TWO: argument,
gauntlet_tchannel.Numberz_THREE: argument,
},
2: {
gauntlet_tchannel.Numberz_SIX: &gauntlet_tchannel.Insanity{},
},
}
return result, nil
}
func (thriftTestHandler) TestMapMap(ctx thrift.Context, hello int32) (map[int32]map[int32]int32, error) {
ctx.SetResponseHeaders(ctx.Headers())
result := map[int32]map[int32]int32{
-4: {
-4: -4,
-3: -3,
-2: -2,
-1: -1,
},
4: {
1: 1,
2: 2,
3: 3,
4: 4,
},
}
return result, nil
}
func (thriftTestHandler) TestMulti(
ctx thrift.Context,
arg0 int8,
arg1 int32,
arg2 int64,
arg3 map[int16]string,
arg4 gauntlet_tchannel.Numberz,
arg5 gauntlet_tchannel.UserId,
) (*gauntlet_tchannel.Xtruct, error) {
ctx.SetResponseHeaders(ctx.Headers())
hello := "Hello2"
result := &gauntlet_tchannel.Xtruct{
StringThing: &hello,
ByteThing: &arg0,
I32Thing: &arg1,
I64Thing: &arg2,
}
return result, nil
}
func (thriftTestHandler) TestList(ctx thrift.Context, thing []int32) ([]int32, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (thriftTestHandler) TestSet(ctx thrift.Context, thing map[int32]bool) (map[int32]bool, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (thriftTestHandler) TestStringMap(ctx thrift.Context, thing map[string]string) (map[string]string, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (thriftTestHandler) TestNest(ctx thrift.Context, thing *gauntlet_tchannel.Xtruct2) (*gauntlet_tchannel.Xtruct2, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func isAdmin(ctx thrift.Context) bool {
return ctx.Headers()["user"] == "admin"
}
func (thriftTestHandler) TestVoid(ctx thrift.Context) error {
ctx.SetResponseHeaders(ctx.Headers())
return nil
}
func (secondServiceHandler) SecondtestString(ctx thrift.Context, thing string) (string, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (secondServiceHandler) BlahBlah(ctx thrift.Context) error {
ctx.SetResponseHeaders(ctx.Headers())
return nil
}
func (thriftTestHandler) TestEnum(ctx thrift.Context, thing gauntlet_tchannel.Numberz) (gauntlet_tchannel.Numberz, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (thriftTestHandler) TestI64(ctx thrift.Context, thing int64) (int64, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (thriftTestHandler) TestDouble(ctx thrift.Context, thing float64) (float64, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
func (h *echoThriftHandler) Echo(ctx thrift.Context, ping *echo.Ping) (*echo.Pong, error) {
ctx.SetResponseHeaders(ctx.Headers())
return &echo.Pong{Boop: ping.Beep}, nil
}
func (thriftTestHandler) TestBinary(ctx thrift.Context, thing []byte) ([]byte, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}
// HasForwardedHeader takes the headers that came in via TChannel and looks for
// the precense of a specific ringpop header to see if ringpop already forwarded
// the message. When a message has already been forwarded by ringpop the
// forwarding logic should not be used to prevent unbound forwarding.
func HasForwardedHeader(ctx thrift.Context) bool {
_, ok := ctx.Headers()[forwardedHeaderName]
return ok
}
func (thriftTestHandler) TestTypedef(ctx thrift.Context, thing gauntlet_tchannel.UserId) (gauntlet_tchannel.UserId, error) {
ctx.SetResponseHeaders(ctx.Headers())
return thing, nil
}