fix: add handler for invalid token requests (#1325)

Signed-off-by: Anish Ramasekar <anish.ramasekar@gmail.com>

Author: aramase

pkg/nmi/server/server.go

@@ -40,6 +40,17 @@ const (
 	headerRetryAfter   = "Retry-After"
 )
 
+var (
+	// invalidTokenPathMatcher matches the token path that is not supported by IMDS
+	// this handler is configured right after the token path handler to block requests with
+	// invalid token path instead of sending it to IMDS.
+	// we don't have to handle case sensitivity for "/identity/" as that's rejected by IMDS
+	invalidTokenPathMatcher = mux.MatcherFunc(func(req *http.Request, rm *mux.RouteMatch) bool {
+		r := regexp.MustCompile("/(?i:metadata)/identity(.*?)oauth2(.*?)token") // #nosec
+		return r.MatchString(req.URL.Path)
+	})
+)
+
 // Server encapsulates all of the parameters necessary for starting up
 // the server. These can be set via command line.
 type Server struct {
@@ -98,7 +109,15 @@ func (s *Server) Run() error {
 	go s.updateIPTableRules()
 
 	rtr := mux.NewRouter()
+	// Flow for the request is as follows:
+	// 1. If the request is for token, then it will be handled by tokenHandler post validation.
+	// 2. If the request is for token but the path is invalid, then it will be handled by invalidTokenPathHandler.
+	// 3. If the request is for host token, then it will be handled by hostTokenHandler.
+	// 4. If the request is for instance metadata
+	//    4.1 If blockInstanceMetadata is set to true, then it will be handled by blockInstanceMetadataHandler (deny access to instance metadata).
+	// 5. If the request is for any other path, it will be proxied to IMDS and the response will be returned to the caller.
 	rtr.PathPrefix(tokenPathPrefix).Handler(appHandler(s.msiHandler))
+	rtr.MatcherFunc(invalidTokenPathMatcher).HandlerFunc(invalidTokenPathHandler)
 	rtr.PathPrefix(hostTokenPathPrefix).Handler(appHandler(s.hostHandler))
 	if s.BlockInstanceMetadata {
 		rtr.PathPrefix(instancePathPrefix).HandlerFunc(forbiddenHandler)
@@ -581,6 +600,11 @@ func forbiddenHandler(w http.ResponseWriter, r *http.Request) {
 	http.Error(w, "Request blocked by AAD Pod Identity NMI", http.StatusForbidden)
 }
 
+// invalidTokenPathHandler responds to invalid token requests with HTTP 400 Bad Request
+func invalidTokenPathHandler(w http.ResponseWriter, r *http.Request) {
+	http.Error(w, "Invalid request", http.StatusBadRequest)
+}
+
 func copyHeader(dst, src http.Header) {
 	for k, vv := range src {
 		for _, v := range vv {

pkg/nmi/server/server_test.go

@@ -217,6 +217,36 @@ func TestRouterPathPrefix(t *testing.T) {
 			expectedStatusCode: http.StatusOK,
 			expectedBody:       "default_handler",
 		},
+		{
+			name:               "invalid token request with \\oauth2",
+			url:                `/metadata/identity\oauth2/token/`,
+			expectedStatusCode: http.StatusOK,
+			expectedBody:       "invalid_request_handler",
+		},
+		{
+			name:               "invalid token request with \\token",
+			url:                `/metadata/identity/oauth2\token/`,
+			expectedStatusCode: http.StatusOK,
+			expectedBody:       "invalid_request_handler",
+		},
+		{
+			name:               "invalid token request with \\oauth2\\token",
+			url:                `/metadata/identity\oauth2\token/`,
+			expectedStatusCode: http.StatusOK,
+			expectedBody:       "invalid_request_handler",
+		},
+		{
+			name:               "invalid token request with mix of / and \\",
+			url:                `/metadata/identity/\oauth2\token/`,
+			expectedStatusCode: http.StatusOK,
+			expectedBody:       "invalid_request_handler",
+		},
+		{
+			name:               "invalid token request with multiple \\",
+			url:                `/metadata/identity\\\oauth2\\token/`,
+			expectedStatusCode: http.StatusOK,
+			expectedBody:       "invalid_request_handler",
+		},
 	}
 
 	for _, test := range tests {
@@ -225,6 +255,7 @@ func TestRouterPathPrefix(t *testing.T) {
 			defer teardown()
 
 			rtr.PathPrefix(tokenPathPrefix).HandlerFunc(testTokenHandler)
+			rtr.MatcherFunc(invalidTokenPathMatcher).HandlerFunc(testInvalidRequestHandler)
 			rtr.PathPrefix(hostTokenPathPrefix).HandlerFunc(testHostTokenHandler)
 			rtr.PathPrefix(instancePathPrefix).HandlerFunc(testInstanceHandler)
 			rtr.PathPrefix("/").HandlerFunc(testDefaultHandler)
@@ -263,3 +294,7 @@ func testInstanceHandler(w http.ResponseWriter, r *http.Request) {
 func testDefaultHandler(w http.ResponseWriter, r *http.Request) {
 	fmt.Fprintf(w, "default_handler\n")
 }
+
+func testInvalidRequestHandler(w http.ResponseWriter, r *http.Request) {
+	fmt.Fprintf(w, "invalid_request_handler\n")
+}