schnutibox/vendor/github.com/twitchtv/twirp/context.go
Marvin Steadfast ae35d9ab41
Some checks failed
continuous-integration/drone/push Build is failing
uses buf for compiling proto files and implements the server service
2021-05-05 11:14:17 +02:00

143 lines
5.2 KiB
Go

// Copyright 2018 Twitch Interactive, Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License"). You may not
// use this file except in compliance with the License. A copy of the License is
// located at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// or in the "license" file accompanying this file. This file is distributed on
// an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.
package twirp
import (
"context"
"errors"
"net/http"
"github.com/twitchtv/twirp/internal/contextkeys"
)
// MethodName extracts the name of the method being handled in the given
// context. If it is not known, it returns ("", false).
func MethodName(ctx context.Context) (string, bool) {
name, ok := ctx.Value(contextkeys.MethodNameKey).(string)
return name, ok
}
// ServiceName extracts the name of the service handling the given context. If
// it is not known, it returns ("", false).
func ServiceName(ctx context.Context) (string, bool) {
name, ok := ctx.Value(contextkeys.ServiceNameKey).(string)
return name, ok
}
// PackageName extracts the fully-qualified protobuf package name of the service
// handling the given context. If it is not known, it returns ("", false). If
// the service comes from a proto file that does not declare a package name, it
// returns ("", true).
//
// Note that the protobuf package name can be very different than the go package
// name; the two are unrelated.
func PackageName(ctx context.Context) (string, bool) {
name, ok := ctx.Value(contextkeys.PackageNameKey).(string)
return name, ok
}
// StatusCode retrieves the status code of the response (as string like "200").
// If it is known returns (status, true).
// If it is not known, it returns ("", false).
func StatusCode(ctx context.Context) (string, bool) {
code, ok := ctx.Value(contextkeys.StatusCodeKey).(string)
return code, ok
}
// WithHTTPRequestHeaders stores an http.Header in a context.Context. When
// using a Twirp-generated client, you can pass the returned context
// into any of the request methods, and the stored header will be
// included in outbound HTTP requests.
//
// This can be used to set custom HTTP headers like authorization tokens or
// client IDs. But note that HTTP headers are a Twirp implementation detail,
// only visible by middleware, not by the server implementation.
//
// WithHTTPRequestHeaders returns an error if the provided http.Header
// would overwrite a header that is needed by Twirp, like "Content-Type".
func WithHTTPRequestHeaders(ctx context.Context, h http.Header) (context.Context, error) {
if _, ok := h["Accept"]; ok {
return nil, errors.New("provided header cannot set Accept")
}
if _, ok := h["Content-Type"]; ok {
return nil, errors.New("provided header cannot set Content-Type")
}
if _, ok := h["Twirp-Version"]; ok {
return nil, errors.New("provided header cannot set Twirp-Version")
}
copied := make(http.Header, len(h))
for k, vv := range h {
if vv == nil {
copied[k] = nil
continue
}
copied[k] = make([]string, len(vv))
copy(copied[k], vv)
}
return context.WithValue(ctx, contextkeys.RequestHeaderKey, copied), nil
}
func HTTPRequestHeaders(ctx context.Context) (http.Header, bool) {
h, ok := ctx.Value(contextkeys.RequestHeaderKey).(http.Header)
return h, ok
}
// SetHTTPResponseHeader sets an HTTP header key-value pair using a context
// provided by a twirp-generated server, or a child of that context.
// The server will include the header in its response for that request context.
//
// This can be used to respond with custom HTTP headers like "Cache-Control".
// But note that HTTP headers are a Twirp implementation detail,
// only visible by middleware, not by the clients or their responses.
//
// The header will be ignored (noop) if the context is invalid (i.e. using a new
// context.Background() instead of passing the context from the handler).
//
// If called multiple times with the same key, it replaces any existing values
// associated with that key.
//
// SetHTTPResponseHeader returns an error if the provided header key
// would overwrite a header that is needed by Twirp, like "Content-Type".
func SetHTTPResponseHeader(ctx context.Context, key, value string) error {
if key == "Content-Type" {
return errors.New("header key can not be Content-Type")
}
responseWriter, ok := ctx.Value(contextkeys.ResponseWriterKey).(http.ResponseWriter)
if ok {
responseWriter.Header().Set(key, value)
} // invalid context is ignored, not an error, this is to allow easy unit testing with mock servers
return nil
}
// AddHTTPResponseHeader behaves like SetHTTPResponseHeader,
// but it appends the key-value pair to the header (instead of replacing it).
//
// AddHTTPResponseHeader returns an error if the key is "Content-Type".
func AddHTTPResponseHeader(ctx context.Context, key, value string) error {
if key == "Content-Type" {
return errors.New("header key can not be Content-Type")
}
responseWriter, ok := ctx.Value(contextkeys.ResponseWriterKey).(http.ResponseWriter)
if ok {
responseWriter.Header().Add(key, value)
} // invalid context is ignored, not an error, this is to allow easy unit testing with mock servers
return nil
}