oauth2/router.go

package oauth2

import (
	"somehole.com/common/log"
	"somehole.com/service/oauth2/server"
	"somehole.com/service/router"
)

type Endpoint struct {
	CallbackEndpoint string
	TokenEndpoint    string
}

var endpointDefault = Endpoint{
	CallbackEndpoint: "/callback",
	TokenEndpoint:    "/token",
}

type Router struct {
	*router.Router
	log.Logger
	Endpoint
}

func NewRouter(r *router.Router, logger log.Logger, endpointOverride ...Endpoint) (ro *Router) {
	var endpoint Endpoint
	if r == nil {
		r = router.NewRouter(nil, nil)
	}
	if len(endpointOverride) == 1 {
		endpoint = endpointOverride[0]
	} else {
		endpoint = endpointDefault
	}
	r.AddRequiredRoutes([]string{
		endpoint.CallbackEndpoint,
		endpoint.TokenEndpoint,
	})
	if logger == nil {
		logger = log.NewPlainLogger()
	}
	return &Router{
		Router:   r,
		Logger:   logger,
		Endpoint: endpoint,
	}
}

func (ro *Router) RegisterCallbackServer(srv server.CallbackServer) *Router {
	ro.Register(ro.CallbackEndpoint, router.NewServer(srv.Callback).SetLogger(ro.Logger))
	return ro
}

func (ro *Router) RegisterTokenServer(srv server.TokenServer) *Router {
	ro.Register(ro.TokenEndpoint, router.NewServer(srv.Token).SetLogger(ro.Logger))
	return ro
}