code.vegaprotocol.io/vega@v0.79.0/core/api/rest/middleware.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package rest
    17  
    18  import (
    19  	"net/http"
    20  	"time"
    21  
    22  	"code.vegaprotocol.io/vega/core/metrics"
    23  	vgcontext "code.vegaprotocol.io/vega/libs/context"
    24  	vfmt "code.vegaprotocol.io/vega/libs/fmt"
    25  	vghttp "code.vegaprotocol.io/vega/libs/http"
    26  	"code.vegaprotocol.io/vega/logging"
    27  )
    28  
    29  // RemoteAddrMiddleware is a middleware adding to the current request context the
    30  // address of the caller.
    31  func RemoteAddrMiddleware(log *logging.Logger, next http.Handler) http.Handler {
    32  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    33  		ip, err := vghttp.RemoteAddr(r)
    34  		if err != nil {
    35  			log.Debug("Failed to get remote address in middleware",
    36  				logging.String("remote-addr", r.RemoteAddr),
    37  				logging.String("x-forwarded-for", vfmt.Escape(r.Header.Get("X-Forwarded-For"))),
    38  			)
    39  		} else {
    40  			r = r.WithContext(vgcontext.WithRemoteIPAddr(r.Context(), ip))
    41  		}
    42  		next.ServeHTTP(w, r)
    43  	})
    44  }
    45  
    46  // MetricCollectionMiddleware records the request and the time taken to service it.
    47  func MetricCollectionMiddleware(next http.Handler) http.Handler {
    48  	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    49  		start := time.Now()
    50  		next.ServeHTTP(w, r)
    51  		end := time.Now()
    52  
    53  		// Update the call count and timings in metrics
    54  		timetaken := end.Sub(start)
    55  
    56  		metrics.APIRequestAndTimeREST(r.Method, r.RequestURI, timetaken.Seconds())
    57  	})
    58  }