server.go

package server

/*
Copyright (c) 2018 Vereign AG [https://www.vereign.com]

This is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

import (
	"fmt"
	"log"
	"net"
	"net/http"
	"strings"

	"github.com/grpc-ecosystem/grpc-gateway/runtime"

	"golang.org/x/net/context"

	"code.vereign.com/code/key-storage-agent/handler"
	"code.vereign.com/code/key-storage-agent/session"
	"code.vereign.com/code/viam-apis/authentication"
	"code.vereign.com/code/viam-apis/data-storage-agent/client"
	api "code.vereign.com/code/viam-apis/key-storage-agent/api"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/metadata"
	"github.com/spf13/viper"
)

// private type for Context keys
type contextKey int

const (
	clientIDKey contextKey = iota
)

var pkgCertFile string
var pkgKeyFile string
var pkgCaCertFile string

func credMatcher(headerName string) (mdName string, ok bool) {
	if headerName == "Session" {
		return headerName, true
	}
	return "", false
}

// authenticateAgent check the client credentials
func authenticateClient(ctx context.Context, s *handler.KeyStorageServerImpl, invokedMethod string) (string, error) {
	if md, ok := metadata.FromIncomingContext(ctx); ok {

		clientAuth := &authentication.Authentication{
			Uuid:    strings.Join(md["uuid"], ""),
			Session: strings.Join(md["session"], ""),
		}

		viamAuth := &authentication.Authentication{
			Uuid:    viper.GetString("viamUUID"),
			Session: viper.GetString("viamSession"),
		}

		sessionClient := &client.DataStorageClientImpl{}
		sessionClient.SetUpClient(viamAuth, viper.GetString("dataStorageUrl"), pkgCertFile, pkgKeyFile, pkgCaCertFile, viper.GetInt("maxMessageSize"))
		defer sessionClient.CloseClient()

		if clientAuth.Uuid == viamAuth.Uuid {
			if clientAuth.Session != viamAuth.Session {
				return "", fmt.Errorf("bad session %s", clientAuth.Session)
			}
		} else {
			if session.CheckSession(clientAuth.Uuid, clientAuth.Session, sessionClient) == false {
				return "", fmt.Errorf("bad session %s", clientAuth.Session)
			}
		}

		log.Printf("authenticated uuid: %s", clientAuth.Uuid)

		return clientAuth.Uuid, nil
	}
	return "", fmt.Errorf("missing credentials")
}

// unaryInterceptor call authenticateClient with current context
func unaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler1 grpc.UnaryHandler) (interface{}, error) {
	s, ok := info.Server.(*handler.KeyStorageServerImpl)
	log.Println("Invoked method: " + info.FullMethod)
	if !ok {
		return nil, fmt.Errorf("unable to cast server")
	}
	clientID, err := authenticateClient(ctx, s, info.FullMethod)
	if err != nil {
		log.Printf("Error: %v", err)
		return nil, err
	}

	ctx = context.WithValue(ctx, clientIDKey, clientID)

	return handler1(ctx, req)
}

func StartGRPCServer(address, certFilePath, privateKeyFilePath, caCertFilePath, vereignCertFilePath, vereignPrivateKeyFilePath, dataStorageAddress string, maxMessageSize int) error {
	pkgCertFile = certFilePath
	pkgKeyFile = privateKeyFilePath
	pkgCaCertFile = caCertFilePath

	// create a listener on TCP port
	lis, err := net.Listen("tcp", address)
	if err != nil {
		log.Printf("Error: %v", err)
		return fmt.Errorf("failed to listen: %v", err)
	}

	// create a server instance
	s := handler.KeyStorageServerImpl{
		DataStorageUrl:            dataStorageAddress,
		CertFilePath:              certFilePath,
		KeyFilePath:			   privateKeyFilePath,
		CaCertFilePath:			   caCertFilePath,
		VereignCertFilePath:       vereignCertFilePath,
		VereignPrivateKeyFilePath: vereignPrivateKeyFilePath,
		MaxMessageSize:			   maxMessageSize,
	}

	// Create the TLS credentials
	creds, err := credentials.NewServerTLSFromFile(certFilePath, privateKeyFilePath)
	if err != nil {
		log.Printf("Error: %v", err)
		return fmt.Errorf("could not load TLS keys: %s", err)
	}

	// Create an array of gRPC options with the credentials
	opts := []grpc.ServerOption{
		grpc.Creds(creds),
		grpc.UnaryInterceptor(unaryInterceptor), 
		grpc.MaxRecvMsgSize(viper.GetInt("maxMessageSize")*1024*1024),
	}

	// create a gRPC server object
	grpcServer := grpc.NewServer(opts...)

	// attach the CalcMinimumDistance service to the server
	api.RegisterKeyStorageServer(grpcServer, &s)

	// start the server
	log.Printf("starting HTTP/2 gRPC server on %s", address)
	if err := grpcServer.Serve(lis); err != nil {
		return fmt.Errorf("failed to serve: %s", err)
	}

	return nil
}

func StartRESTServer(address, grpcAddress, certFile string) error {
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	defer cancel()

	mux := runtime.NewServeMux(runtime.WithIncomingHeaderMatcher(credMatcher))

	creds, err := credentials.NewClientTLSFromFile(certFile, "")
	if err != nil {
		log.Printf("Error: %v", err)
		return fmt.Errorf("could not load TLS certificate: %s", err)
	}

	// Setup the client gRPC options
	opts := []grpc.DialOption{grpc.WithTransportCredentials(creds)}

	// Register RedisStorageServer
	err = api.RegisterKeyStorageHandlerFromEndpoint(ctx, mux, grpcAddress, opts)
	if err != nil {
		log.Printf("Error: %v", err)
		return fmt.Errorf("could not register service RedisStorageServer: %s", err)
	}

	log.Printf("starting HTTP/1.1 REST server on %s", address)
	http.ListenAndServe(address, mux)

	return nil
}