github.com/status-im/status-go@v1.1.0/server/pairing/challenge_management.go

package pairing

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"

	"github.com/btcsuite/btcutil/base58"
	"github.com/gorilla/sessions"
	"go.uber.org/zap"
)

const (
	// Session names
	sessionChallenge = "challenge"
	sessionBlocked   = "blocked"
)

type ChallengeError struct {
	Text     string
	HTTPCode int
}

func (ce *ChallengeError) Error() string {
	return fmt.Sprintf("%s : %d", ce.Text, ce.HTTPCode)
}

func makeCookieStore() (*sessions.CookieStore, error) {
	auth := make([]byte, 64)
	_, err := rand.Read(auth)
	if err != nil {
		return nil, err
	}

	enc := make([]byte, 32)
	_, err = rand.Read(enc)
	if err != nil {
		return nil, err
	}

	return sessions.NewCookieStore(auth, enc), nil
}

// ChallengeGiver is responsible for generating challenges and checking challenge responses
type ChallengeGiver struct {
	cookieStore *sessions.CookieStore
	encryptor   *PayloadEncryptor
	logger      *zap.Logger
	authedIP    net.IP
}

func NewChallengeGiver(e *PayloadEncryptor, logger *zap.Logger) (*ChallengeGiver, error) {
	cs, err := makeCookieStore()
	if err != nil {
		return nil, err
	}

	return &ChallengeGiver{
		cookieStore: cs,
		encryptor:   e.Renew(),
		logger:      logger,
	}, nil
}

func (cg *ChallengeGiver) getIP(r *http.Request) (net.IP, error) {
	h, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		cg.logger.Error("getIP: h, _, err := net.SplitHostPort(r.RemoteAddr)", zap.Error(err), zap.String("r.RemoteAddr", r.RemoteAddr))
		return nil, &ChallengeError{"error", http.StatusInternalServerError}
	}
	return net.ParseIP(h), nil
}

func (cg *ChallengeGiver) registerClientIP(r *http.Request) error {
	IP, err := cg.getIP(r)
	if err != nil {
		return err
	}
	cg.authedIP = IP
	return nil
}

func (cg *ChallengeGiver) validateClientIP(r *http.Request) error {
	// If we haven't registered yet register the IP
	if cg.authedIP == nil || len(cg.authedIP) == 0 {
		err := cg.registerClientIP(r)
		if err != nil {
			return err
		}
	}

	// Then compare the current req RemoteIP with the authed IP
	IP, err := cg.getIP(r)
	if err != nil {
		return err
	}

	if !cg.authedIP.Equal(IP) {
		cg.logger.Error(
			"request RemoteAddr does not match authedIP: expected '%s', received '%s'",
			zap.String("expected", cg.authedIP.String()),
			zap.String("received", IP.String()),
		)
		return &ChallengeError{"forbidden", http.StatusForbidden}
	}
	return nil
}

func (cg *ChallengeGiver) getSession(r *http.Request) (*sessions.Session, error) {
	s, err := cg.cookieStore.Get(r, sessionChallenge)
	if err != nil {
		cg.logger.Error("checkChallengeResponse: cg.cookieStore.Get(r, sessionChallenge)", zap.Error(err), zap.String("sessionChallenge", sessionChallenge))
		return nil, &ChallengeError{"error", http.StatusInternalServerError}
	}
	return s, nil
}

func (cg *ChallengeGiver) generateNewChallenge(s *sessions.Session, w http.ResponseWriter, r *http.Request) ([]byte, error) {
	challenge := make([]byte, 64)
	_, err := rand.Read(challenge)
	if err != nil {
		cg.logger.Error("regenerateNewChallenge: _, err = rand.Read(challenge)", zap.Error(err))
		return nil, &ChallengeError{"error", http.StatusInternalServerError}
	}

	s.Values[sessionChallenge] = challenge
	err = s.Save(r, w)
	if err != nil {
		cg.logger.Error("regenerateNewChallenge: err = s.Save(r, w)", zap.Error(err))
		return nil, &ChallengeError{"error", http.StatusInternalServerError}
	}

	return challenge, nil
}

func (cg *ChallengeGiver) block(s *sessions.Session, w http.ResponseWriter, r *http.Request) error {
	s.Values[sessionBlocked] = true
	err := s.Save(r, w)
	if err != nil {
		cg.logger.Error("block: err = s.Save(r, w)", zap.Error(err))
		return &ChallengeError{"error", http.StatusInternalServerError}
	}

	return &ChallengeError{"forbidden", http.StatusForbidden}
}

func (cg *ChallengeGiver) checkChallengeResponse(w http.ResponseWriter, r *http.Request) error {
	err := cg.validateClientIP(r)
	if err != nil {
		return err
	}

	s, err := cg.getSession(r)
	if err != nil {
		return err
	}

	blocked, ok := s.Values[sessionBlocked].(bool)
	if ok && blocked {
		return &ChallengeError{"forbidden", http.StatusForbidden}
	}

	// If the request header doesn't include a challenge don't punish the client, just throw a 403
	clientChallengeResp := r.Header.Get(sessionChallenge)
	if clientChallengeResp == "" {
		return &ChallengeError{"forbidden", http.StatusForbidden}
	}

	dcr, err := cg.encryptor.decryptPlain(base58.Decode(clientChallengeResp))
	if err != nil {
		cg.logger.Error("checkChallengeResponse: cg.encryptor.decryptPlain(base58.Decode(clientChallengeResp))", zap.Error(err), zap.String("clientChallengeResp", clientChallengeResp))
		return &ChallengeError{"error", http.StatusInternalServerError}
	}

	// If the challenge is not in the session store don't punish the client, just throw a 403
	challenge, ok := s.Values[sessionChallenge].([]byte)
	if !ok {
		return &ChallengeError{"forbidden", http.StatusForbidden}
	}

	// Only if we have both a challenge in the session store and in the request header
	// do we entertain blocking the client. Because then we know someone is trying to be sneaky.
	if !bytes.Equal(dcr, challenge) {
		return cg.block(s, w, r)
	}

	// If every is ok, generate a new challenge for the next req
	_, err = cg.generateNewChallenge(s, w, r)
	return err
}

func (cg *ChallengeGiver) getChallenge(w http.ResponseWriter, r *http.Request) ([]byte, error) {
	err := cg.validateClientIP(r)
	if err != nil {
		return nil, err
	}

	s, err := cg.getSession(r)
	if err != nil {
		return nil, err
	}

	challenge, ok := s.Values[sessionChallenge].([]byte)
	if !ok {
		challenge, err = cg.generateNewChallenge(s, w, r)
		if err != nil {
			return nil, err
		}
	}
	return challenge, nil
}

// ChallengeTaker is responsible for storing and performing server challenges
type ChallengeTaker struct {
	encryptor       *PayloadEncryptor
	serverChallenge []byte
}

func NewChallengeTaker(e *PayloadEncryptor) *ChallengeTaker {
	return &ChallengeTaker{
		encryptor: e.Renew(),
	}
}

func (ct *ChallengeTaker) SetChallenge(resp *http.Response) error {
	challenge, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return err
	}
	ct.serverChallenge = challenge
	return nil
}

func (ct *ChallengeTaker) DoChallenge(req *http.Request) error {
	if ct.serverChallenge != nil {
		ec, err := ct.encryptor.encryptPlain(ct.serverChallenge)
		if err != nil {
			return err
		}

		req.Header.Set(sessionChallenge, base58.Encode(ec))
	}
	return nil
}