Skip to content

Instantly share code, notes, and snippets.

@kvii
Created March 13, 2024 02:03
Show Gist options
  • Save kvii/94cf3a1ef9d7cc09f2417781f2e4f8a9 to your computer and use it in GitHub Desktop.
Save kvii/94cf3a1ef9d7cc09f2417781f2e4f8a9 to your computer and use it in GitHub Desktop.
Custom http status code in kratos.
package main
import (
"context"
"log"
"github.com/go-kratos/kratos/v2"
"github.com/go-kratos/kratos/v2/transport/http"
)
type customStatusCode struct{}
var customStatusCodeKey = new(customStatusCode)
// mix a custom code value to ctx
func WithCustomStatusCode(ctx context.Context, code int) context.Context {
return context.WithValue(ctx, customStatusCodeKey, code)
}
// get value from ctx
func FromCustomStatusCode(ctx context.Context) (int, bool) {
code, ok := ctx.Value(customStatusCodeKey).(int)
return code, ok
}
func WithHttpCode(ctx context.Context, code int) {
if hc, ok := ctx.(http.Context); ok {
r := hc.Request()
ctx := r.Context()
ctx = WithCustomStatusCode(ctx, code)
hc.Reset(hc.Response(), r.WithContext(ctx))
}
}
func main() {
srv := http.NewServer(
http.Address(":8000"),
http.ResponseEncoder(func(w http.ResponseWriter, r *http.Request, i any) error {
if code, ok := FromCustomStatusCode(r.Context()); ok {
w.WriteHeader(code)
}
return http.DefaultResponseEncoder(w, r, i)
}),
)
// your service method
h := func(ctx context.Context) (any, error) {
WithHttpCode(ctx, 201)
reply := map[string]any{
"id": 1,
"name": "a",
}
return reply, nil
}
// simulate what happened in generated xx_http.pb.go
srv.Route("/").GET("/a", func(ctx http.Context) error {
reply, _ := h(ctx)
return ctx.Result(200, reply)
})
k := kratos.New(kratos.Server(srv))
if err := k.Run(); err != nil {
log.Fatal(err)
}
// $ curl http://localhost:8000/a -v
//
// ...
// < HTTP/1.1 201 Created
// < Content-Type: application/json
// ...
// {"id":1,"name":"a"}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment