示例#1
0
文件: cluster.go 项目: naunga/vault
// WrapHandlerForClustering takes in Vault's HTTP handler and returns a setup
// function that returns both the original handler and one wrapped with cluster
// methods
func WrapHandlerForClustering(handler http.Handler, logger log.Logger) func() (http.Handler, http.Handler) {
	return func() (http.Handler, http.Handler) {
		// This mux handles cluster functions (right now, only forwarded requests)
		mux := http.NewServeMux()
		mux.HandleFunc("/cluster/local/forwarded-request", func(w http.ResponseWriter, req *http.Request) {
			freq, err := forwarding.ParseForwardedHTTPRequest(req)
			if err != nil {
				if logger != nil {
					logger.Error("http/forwarded-request-server: error parsing forwarded request", "error", err)
				}

				w.Header().Add("Content-Type", "application/json")

				// The response writer here is different from
				// the one set in Vault's HTTP handler.
				// Hence, set the Cache-Control explicitly.
				w.Header().Set("Cache-Control", "no-store")

				w.WriteHeader(http.StatusInternalServerError)

				type errorResponse struct {
					Errors []string
				}
				resp := &errorResponse{
					Errors: []string{
						err.Error(),
					},
				}

				enc := json.NewEncoder(w)
				enc.Encode(resp)
				return
			}

			// To avoid the risk of a forward loop in some pathological condition,
			// set the no-forward header
			freq.Header.Set(IntNoForwardingHeaderName, "true")
			handler.ServeHTTP(w, freq)
		})

		return handler, mux
	}
}
示例#2
0
func acceptStreams(logger log.Logger, session *yamux.Session, streamCh chan net.Conn) grim.TaskFunc {
	return func(ctx context.Context) {
		defer close(streamCh)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				stream, err := session.Accept()
				if err != nil {
					if err != io.EOF {
						logger.Error("multiplex conn accept failed", "err", err)
					}
					return
				}
				streamCh <- stream
			}
		}
	}
}