Skip to content

Commit

Permalink
[rest_gateway] Add authentication to gRPC REST gateway and include un…
Browse files Browse the repository at this point in the history
…it testing (#1453)

Features and improvements

1. JWT authentication for gRPC REST gateway
- Integrated JWT authentication for secure communication.
- Implemented middleware to validate JWT tokens from HTTP requests'
authorization headers.
2. Include unit testing
- Added Go tests for the REST gateway to ensure authentication
middleware functionality and error handling.

Co-authored-by: Zachary Fong <[email protected]>
Co-authored-by: Ramon Figueiredo <[email protected]>
Co-authored-by: Diego Tavares <[email protected]>
  • Loading branch information
4 people authored Aug 9, 2024
1 parent 1219377 commit db5a805
Show file tree
Hide file tree
Showing 6 changed files with 212 additions and 48 deletions.
1 change: 1 addition & 0 deletions rest_gateway/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
.env
40 changes: 13 additions & 27 deletions rest_gateway/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,30 +1,9 @@
# Build from project root
FROM rockylinux:9.3 as rocky-golang
FROM centos7.6-go1.21:latest AS build

WORKDIR /src/go
RUN dnf install -y 'dnf-command(config-manager)' && dnf config-manager --set-enabled crb

# WARN: Download do tarball and update file path accordingly
COPY rest_gateway/Go_1.22.2_Linux_ARM64.tar.gz golang.tar.gz

RUN tar -C /usr/local -xzf golang.tar.gz
ENV PATH="$PATH:/usr/local/go/bin"

# WARN: Uncoment if your environment requires a proxy
#ENV GOPROXY=artifactory.yourcompany.com/go-proxy
ENV GO111MODULE=on
ENV GOSUMDB=off

RUN go version

RUN rm -rf /src/go

FROM rocky-golang AS build

RUN dnf install -y \
RUN yum install -y \
git \
protobuf-compiler \
&& dnf clean all
protobuf3-compiler \
&& yum clean all
WORKDIR /app
ENV PATH=$PATH:/root/go/bin:/opt/protobuf3/usr/bin/

Expand All @@ -33,13 +12,14 @@ COPY ./proto /app/proto
COPY ./rest_gateway/opencue_gateway /app/opencue_gateway
# COPY ./lib /app/lib


WORKDIR /app/opencue_gateway
RUN go mod init opencue_gateway && go mod tidy

RUN go install \
github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-grpc-gateway \
github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2 \
github.com/golang-jwt/jwt/v5 \
google.golang.org/protobuf/cmd/protoc-gen-go \
google.golang.org/grpc/cmd/protoc-gen-go-grpc

Expand All @@ -58,11 +38,17 @@ RUN protoc -I ../proto/ --grpc-gateway_out ./gen/go \
--grpc-gateway_opt generate_unbound_methods=true \
../proto/*.proto

# Uncomment this to run go tests
# RUN go test -v

# Build project
RUN go build -o grpc_gateway main.go

FROM rockylinux:9.3
FROM centos-7.6.1810:latest
COPY --from=build /app/opencue_gateway/grpc_gateway /app/

# Ensure logs folder is created and has correct permissions
RUN mkdir -p /logs && chmod 755 /logs

EXPOSE 8448
ENTRYPOINT ["/app/grpc_gateway"]
39 changes: 25 additions & 14 deletions rest_gateway/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -6,19 +6,20 @@ A gateway to provide a REST endpoint to opencue gRPC API.

This is a go serviced based on the official [grpc-gateway project](https://github.com/grpc-ecosystem/grpc-gateway)
that compiles opencue's proto files into a go service that provides a REST interface and redirect calls to the
grpc endpoint.
grpc endpoint. All API calls over the REST interface requires an authentication header with a json web token as the bearer.

## Running the service

Running the service is very simple:
* Read and modify the rest_gateway/Dockerfile according to your environment and build the gateway image using docker.
* Run the image providing the environment variable `CUEBOT_ENDPOINT=your.cuebot.server:8443`
**Note:** In the examples below, the REST gateway is available at OPENCUE_REST_GATEWAY_URL. Remember to replace OPENCUE_REST_GATEWAY_URL with the appropriate URL.

## REST interface

All service rpc calls are accessible:
* HTTP method is POST
* URI path is built from the service’s name and method: /<fully qualified service name>/<method name> (e.g.: /show.ShowInterface/FindShow)
* HTTP header must have an authorization with a jwt token as the bearer. e.g:
```headers: {
"Authorization": `Bearer ${jwtToken}`,
},
```
* HTTP body is a JSON with the request object: e.g.:
```proto
message ShowFindShowRequest {
Expand Down Expand Up @@ -61,9 +62,9 @@ message Show {
float default_max_gpus = 11;
}
```
request (gateway running on `http://opencue-gateway.apps.com`):
request (gateway running on `OPENCUE_REST_GATEWAY_URL`):
```bash
curl -i -X POST http://opencue-gateway.apps.com/show.ShowInterface/FindShow -d '{"name": "ashow"}`
curl -i -H "Authorization: Bearer jwtToken" -X POST OPENCUE_REST_GATEWAY_URL/show.ShowInterface/FindShow -d '{"name": "testshow"}`
```
response
```bash
Expand All @@ -74,7 +75,7 @@ Grpc-Metadata-Grpc-Accept-Encoding: gzip
Date: Tue, 12 Dec 2023 18:05:18 GMT
Content-Length: 501
{"show":{"id":"00000000-0000-0000-0000-99999999999999","name":"ashow","defaultMinCores":1,"defaultMaxCores":10,"commentEmail":"middle-tier@imageworks.com","bookingEnabled":true,"dispatchEnabled":true,"active":true,"showStats":{"runningFrames":75,"deadFrames":14,"pendingFrames":1814,"pendingJobs":175,"createdJobCount":"2353643","createdFrameCount":"10344702","renderedFrameCount":"9733366","failedFrameCount":"1096394","reservedCores":252,"reservedGpus":0},"defaultMinGpus":100,"defaultMaxGpus":100000}}
{"show":{"id":"00000000-0000-0000-0000-000000000000","name":"testshow","defaultMinCores":1,"defaultMaxCores":10,"commentEmail":"middle-tier@company.com","bookingEnabled":true,"dispatchEnabled":true,"active":true,"showStats":{"runningFrames":75,"deadFrames":14,"pendingFrames":1814,"pendingJobs":175,"createdJobCount":"2353643","createdFrameCount":"10344702","renderedFrameCount":"9733366","failedFrameCount":"1096394","reservedCores":252,"reservedGpus":0},"defaultMinGpus":100,"defaultMaxGpus":100000}}
```
### Example (getting frames for a job):
Expand Down Expand Up @@ -143,11 +144,11 @@ message FrameSeq {
repeated Frame frames = 1;
}
```
request (gateway running on `http://opencue-gateway.apps.com`):
request (gateway running on `OPENCUE_REST_GATEWAY_URL`):
Note: it is important to include 'page' and 'limit' when getting frames for a job.
```bash
curl -i -X POST http://opencue-gateway.apps.com/job.JobInterface/GetFrames -d '{"job":{"id":"9999999999-b8d7-9999-a29c-99999999999999"}, "req": {"include_finished":true,"page":1,"limit":100}}'
curl -i -H "Authorization: Bearer jwtToken" -X POST OPENCUE_REST_GATEWAY_URL/job.JobInterface/GetFrames -d '{"job":{"id":"00000000-0000-0000-0000-000000000001", "req": {"include_finished":true,"page":1,"limit":100}}'
```
response
```bash
Expand All @@ -157,7 +158,17 @@ grpc-metadata-content-type: application/grpc
grpc-metadata-grpc-accept-encoding: gzip
date: Tue, 13 Feb 2024 17:15:49 GMT
transfer-encoding: chunked
set-cookie: 3d3a38cc45d028e42e93031e0ccc9b1e=534d34fde72242856a7fdadc27260929; path=/; HttpOnly
set-cookie: 01234567890123456789012345678901234567890123456789012345678901234; path=/; HttpOnly
{"frames":{"frames":[{"id":"00000000-0000-0000-0000-000000000002", "name":"0001-bty_tp_3d_123456", "layerName":"bty_tp_3d_123456", "number":1, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":0, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"00000000-0000-0000-0000-000000000003", "name":"0002-bty_tp_3d_123456", "layerName":"bty_tp_3d_123456", "number":2, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":1, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"00000000-0000-0000-0000-000000000004", "name":"0003-bty_tp_3d_083540", "layerName":"bty_tp_3d_123456", "number":3, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":2, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"00000000-0000-0000-0000-000000000005", "name":"0004-bty_tp_3d_083540", "layerName":"bty_tp_3d_123456", "number":4, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":3, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"00000000-0000-0000-0000-000000000006", "name":"0005-bty_tp_3d_083540", "layerName":"bty_tp_3d_123456", "number":5, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":4, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}]}}
```
## Unit testing and system logs
Unit tests for the gRPC REST gateway can be run by uncommenting `RUN go test -v` in the Dockerfile. Unit tests currently cover the following cases for jwtMiddleware (used for authentication):
- valid tokens
- missing tokens
- invalid tokens
- expired tokens
{"frames":{"frames":[{"id":"9999999", "name":"0001-some_frame_0999990", "layerName":"h", "number":1, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":0, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"10fa17d4-9313-4924-86f5-380c5b2a25d8", "name":"0002-some_frame_0999990", "layerName":"some_frame_0999990", "number":2, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":1, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"8d39c602-0b27-4b1e-a09b-4fa35db40e55", "name":"0003-some_frame_0999990", "layerName":"some_frame_0999990", "number":3, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":2, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"d418a837-e974-4716-9105-296f495bc407", "name":"0004-some_frame_0999990", "layerName":"some_frame_0999990", "number":4, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":3, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}, {"id":"d2113372-99999-4c05-8100-9999999", "name":"0005-some_frame_0999990", "layerName":"some_frame_0999990", "number":5, "state":"WAITING", "retryCount":0, "exitStatus":-1, "dispatchOrder":4, "startTime":0, "stopTime":0, "maxRss":"0", "usedMemory":"0", "reservedMemory":"0", "reservedGpuMemory":"0", "lastResource":"/0.00/0", "checkpointState":"DISABLED", "checkpointCount":0, "totalCoreTime":0, "lluTime":1707842141, "totalGpuTime":0, "maxGpuMemory":"0", "usedGpuMemory":"0", "frameStateDisplayOverride":null}]}}
```
System logs are available in /logs and require mounting to be properly tracked. All Stdout are output to both the console and /logs. Here is an example Docker run command that includes addding an environment file and volume mounting: `docker run --env-file ./rest_gateway/.env -v PATH_TO_REST_GATEWAY/logs:/logs -p 8448:8448 restgateway`.
84 changes: 77 additions & 7 deletions rest_gateway/opencue_gateway/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -2,9 +2,14 @@ package main
import (
"context"
"flag"
"fmt"
"io"
"log"
"net/http"
"os"

"strings"

"github.com/golang-jwt/jwt/v5"
"github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials/insecure"
Expand All @@ -13,16 +18,68 @@ import (
gw "opencue_gateway/gen/go" // Update
)

func getEnv(key, fallback string) string {
func getEnv(key string) string {
// Return the value of the environment variable if it's found
if value, ok := os.LookupEnv(key); ok {
return value
} else {
// If the environment variable is not found, output an error and exit the program
log.Fatal(fmt.Sprintf("Error: environment variable '%v' not found", key))
}
return fallback
return ""
}

// Parse and validate the JWT token string
func validateJWTToken(tokenString string, jwtSecret []byte) (*jwt.Token, error) {
return jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
// Ensure that the token's signing method is HMAC
if _, ok := token.Method.(*jwt.SigningMethodHMAC); !ok {
errorString := fmt.Sprintf("Unexpected signing method: %v", token.Header["alg"])
log.Printf(errorString)
return nil, fmt.Errorf(errorString)
}
// Return the secret key for validation
return jwtSecret, nil
})
}

// Middleware to handle token authorization
func jwtMiddleware(next http.Handler, jwtSecret []byte) http.Handler {
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// Get the authorization header and return 401 if there is no header
authHeader := r.Header.Get("Authorization")
if authHeader == "" {
errorString := "Authorization header required"
log.Printf(errorString)
http.Error(w, errorString, http.StatusUnauthorized)
return
}

// Get the token from the header and validate it
tokenString := strings.TrimPrefix(authHeader, "Bearer ")
token, err := validateJWTToken(tokenString, jwtSecret)
if err!=nil {
errorString := fmt.Sprintf("Token validation error: %v", err)
log.Printf(errorString)
http.Error(w, errorString, http.StatusUnauthorized)
return
}
if !token.Valid {
errorString := "Invalid token"
log.Printf(errorString)
http.Error(w, errorString, http.StatusUnauthorized)
return
}

// If token is valid, pass it to the next handler
next.ServeHTTP(w, r)
})
}

func run() error {
grpcServerEndpoint := getEnv("CUEBOT_ENDPOINT", "opencuetest01.your.test.server:8443")
port := getEnv("REST_PORT", "8448")
grpcServerEndpoint := getEnv("CUEBOT_ENDPOINT")
port := getEnv("REST_PORT")
jwtSecret := []byte(getEnv("JWT_AUTH_SECRET"))

ctx := context.Background()
ctx, cancel := context.WithCancel(ctx)
Expand Down Expand Up @@ -73,13 +130,26 @@ func run() error {
if errProc != nil {
return errProc
}


// Create a new HTTP ServeMux with middleware jwtMiddleware to protect the mux
httpMux := http.NewServeMux()
httpMux.Handle("/", jwtMiddleware(mux, jwtSecret))

// Start HTTP server (and proxy calls to gRPC server endpoint)
return http.ListenAndServe(":" + port, mux)
return http.ListenAndServe(":" + port, httpMux)
}

func main() {
// Set up file to capture all log outputs
f, err := os.OpenFile("/logs/opencue_gateway.log", os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0666)
if err != nil {
log.Fatal(err)
}
// Enable output to both Stdout and the log file
mw := io.MultiWriter(os.Stdout, f)
defer f.Close()
log.SetOutput(mw)

flag.Parse()

if err := run(); err != nil {
Expand Down
88 changes: 88 additions & 0 deletions rest_gateway/opencue_gateway/main_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
package main

import (
"net/http"
"net/http/httptest"
"testing"
"time"

"github.com/golang-jwt/jwt/v5"
"github.com/stretchr/testify/assert"
)

func TestJwtMiddleware(t *testing.T) {
jwtSecret := []byte("test_secret")

// Set up a sample handler to use with the middleware
sampleHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.WriteHeader(http.StatusOK)
w.Write([]byte("OK"))
})

// Create a test server with the middleware
ts := httptest.NewServer(jwtMiddleware(sampleHandler, jwtSecret))
defer ts.Close()

t.Run("Valid Token", func(t *testing.T) {
// Generate a valid token
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"sub": "test_user",
"exp": time.Now().Add(time.Hour).Unix(),
})
tokenString, err := token.SignedString(jwtSecret)
assert.NoError(t, err)

// Create a request with the valid token
req, err := http.NewRequest("GET", ts.URL, nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenString)

// Perform the request
res, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.Equal(t, http.StatusOK, res.StatusCode)
})

t.Run("Missing Token", func(t *testing.T) {
// Create a request without a token
req, err := http.NewRequest("GET", ts.URL, nil)
assert.NoError(t, err)

// Perform the request
res, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.Equal(t, http.StatusUnauthorized, res.StatusCode)
})

t.Run("Invalid Token", func(t *testing.T) {
// Create a request with an invalid token
req, err := http.NewRequest("GET", ts.URL, nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer invalid_token")

// Perform the request
res, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.Equal(t, http.StatusUnauthorized, res.StatusCode)
})

t.Run("Expired Token", func(t *testing.T) {
// Generate an expired token
token := jwt.NewWithClaims(jwt.SigningMethodHS256, jwt.MapClaims{
"sub": "test_user",
"exp": time.Now().Add(-time.Hour).Unix(),
})
tokenString, err := token.SignedString(jwtSecret)
assert.NoError(t, err)

// Create a request with the expired token
req, err := http.NewRequest("GET", ts.URL, nil)
assert.NoError(t, err)
req.Header.Set("Authorization", "Bearer "+tokenString)

// Perform the request
res, err := http.DefaultClient.Do(req)
assert.NoError(t, err)
assert.Equal(t, http.StatusUnauthorized, res.StatusCode)
})
}
8 changes: 8 additions & 0 deletions rest_gateway/opencue_gateway/tools/gateway.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
package tools

import (
_ "github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-grpc-gateway"
_ "github.com/grpc-ecosystem/grpc-gateway/v2/protoc-gen-openapiv2"
_ "google.golang.org/grpc/cmd/protoc-gen-go-grpc"
_ "google.golang.org/protobuf/cmd/protoc-gen-go"
)

0 comments on commit db5a805

Please sign in to comment.