neowolves/auth-server/server.go

package main

import (
	"context"
	"fmt"

	"github.com/nspcc-dev/neo-go/pkg/rpcclient"
	"github.com/nspcc-dev/neo-go/pkg/rpcclient/actor"
	"github.com/nspcc-dev/neo-go/pkg/util"
	"github.com/nspcc-dev/neo-go/pkg/wallet"

	"github.com/go-oauth2/oauth2/v4"
	"github.com/go-oauth2/oauth2/v4/errors"
	"github.com/go-oauth2/oauth2/v4/manage"
	"github.com/go-oauth2/oauth2/v4/models"
	"github.com/go-oauth2/oauth2/v4/server"
	"github.com/go-oauth2/oauth2/v4/store"

	"log"
	"log/slog"
	"net/http"
)

type IBlockchainStorage interface {
	GetByID(ctx context.Context, id string) (oauth2.ClientInfo, error) // read
	Set(clt oauth2.ClientInfo) error                                   // create and update
	Delete(id string) error                                            // delete
	CheckPassword(id string, secret string) (bool, error)              // CheckUser
}

type BlockchainStorage struct {
	contract *Contract
}

type StorageClientInfo struct {
	id     string
	secret string
}

func NewBlockchainStorage(actor Actor, hash util.Uint160) *BlockchainStorage {
	return &BlockchainStorage{contract: New(actor, hash)}
}

func (storage BlockchainStorage) GetByID(ctx context.Context, id string) (oauth2.ClientInfo, error) {
	password, err := storage.contract.GetUser(id)
	if err != nil {
		return nil, err
	}

	return &StorageClientInfo{id: id, secret: password.StringLE()}, nil
}

func (storage BlockchainStorage) Set(clt oauth2.ClientInfo) error {
	password, decodeErr := util.Uint256DecodeStringLE(clt.GetSecret())
	if decodeErr != nil {
		return decodeErr
	}

	_, _, err := storage.contract.UpdateUser(clt.GetID(), password)
	if err != nil {
		_, _, err = storage.contract.CreateUser(clt.GetID(), password)

		if err != nil {
			return errors.New("Something went wrong")
		}
	}

	return nil
}

func (storage BlockchainStorage) Delete(id string) error {
	// should we use hash and ValidUntilBlock?
	_, _, res := storage.contract.DeleteUser(id)
	return res
}

func (storage BlockchainStorage) CheckPassword(id string, secret string) (bool, error) {
	password, decodeErr := util.Uint256DecodeStringLE(secret)
	if decodeErr != nil {
		return false, decodeErr
	}

	_, err := storage.contract.CheckUser(id, password)
	if err != nil {
		return false, err
	}

	return true, nil
}

func (model StorageClientInfo) GetID() string {
	return model.id
}

func (model StorageClientInfo) GetSecret() string {
	return model.secret
}

func (model StorageClientInfo) GetDomain() string {
	return model.GetDomain()
}

func (model StorageClientInfo) IsPublic() bool {
	return model.IsPublic()
}

func (model StorageClientInfo) GetUserID() string {
	return model.GetUserID()
}

type InMemoryClient struct {
	Id       string
	Domain   string
	IsPublic bool
}

var clients map[string]InMemoryClient

func addInMemoryClient(id string, domain string, isPublic bool) error {
	result := InMemoryClient{
		Id:       id,
		Domain:   domain,
		IsPublic: isPublic,
	}
	_, contains := clients[id]
	if contains {
		return fmt.Errorf("client with id %s already exist", id)
	} else {
		clients[id] = result
		return nil
	}
}

func getInMemoryClient(id string) (*InMemoryClient, error) {
	result, contains := clients[id]
	if contains {
		return &result, nil
	} else {
		return nil, fmt.Errorf("client with id %s not found", id)
	}
}

func updateInMemoryClientById(id string, client InMemoryClient) error {
	_, contains := clients[id]
	if !contains || id != client.Id {
		return fmt.Errorf("client with id %s not found", id)
	} else {
		clients[id] = client
		return nil
	}
}

func deleteInMemoryClient(id string) {
	delete(clients, id)
}

func main() {
	manager := manage.NewDefaultManager()
	manager.SetAuthorizeCodeTokenCfg(manage.DefaultAuthorizeCodeTokenCfg)

	// contract integration
	fileWallet, _ := wallet.NewWalletFromFile("somewhere")
	defer fileWallet.Close()

	rpcClient, _ := rpcclient.New(context.Background(), "url", rpcclient.Options{})
	rpcActor, _ := actor.NewSimple(rpcClient, fileWallet.Accounts[0])

	// token memory store
	// todo: Implement blockchain store
	manager.MustTokenStorage(store.NewMemoryTokenStore())

	// client memory store
	// todo: Implement blockchain store
	clientStore := store.NewClientStore()

	manager.MapClientStorage(clientStore)

	srv := server.NewDefaultServer(manager)
	srv.SetAllowGetAccessRequest(true)
	srv.SetClientInfoHandler(server.ClientFormHandler)
	manager.SetRefreshTokenCfg(manage.DefaultRefreshTokenCfg)

	srv.SetInternalErrorHandler(func(err error) (re *errors.Response) {
		log.Println("Internal Error:", err.Error())
		return
	})

	srv.SetResponseErrorHandler(func(re *errors.Response) {
		log.Println("Response Error:", re.Error.Error())
	})

	http.HandleFunc("/login", func(w http.ResponseWriter, r *http.Request) {
		srv.HandleTokenRequest(w, r)

		id := r.Header.Get("client_id")
		_, err := getInMemoryClient(id)
		if err != nil {
			slog.Warn("Client with id " + id + "not found")
			w.WriteHeader(http.StatusBadRequest)
		}
	})

	http.HandleFunc("/register", func(writer http.ResponseWriter, request *http.Request) {
		id := request.Header.Get("client_id")
		secret := request.Header.Get("client_secret")

		// check whether client exists
		_, err := clientStore.GetByID(context.Background(), id)
		if err == nil {
			slog.Warn("Client with id " + id + "already exists.")
			writer.WriteHeader(http.StatusBadRequest)
			return
		}

		// add client's credentials to blockchain
		err = clientStore.Set(id, &models.Client{
			ID:     id,
			Secret: secret,
		})

		if err != nil {
			slog.Error("Failed to register user with id "+id+".", err)
			writer.WriteHeader(http.StatusInternalServerError)
			return
		}

		client, _ := clientStore.GetByID(context.Background(), id)
		err = addInMemoryClient(
			client.GetID(),
			client.GetDomain(),
			client.IsPublic())
		if err != nil {
			slog.Error(err.Error())
			return
		}

		writer.WriteHeader(http.StatusOK)
	})

	// for tests
	http.HandleFunc("/protected", validateToken(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte("Hello, I'm protected"))
	}, srv))

	log.Fatal(http.ListenAndServe(":9096", nil))
}

func validateToken(f http.HandlerFunc, srv *server.Server) http.HandlerFunc {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		_, err := srv.ValidationBearerToken(r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}

		f.ServeHTTP(w, r)
	})
}