/*	$OpenBSD: wg_noise.h,v 1.2 2020/12/09 05:53:33 tb Exp $ */
/*
 * Copyright (C) 2015-2020 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 * Copyright (C) 2019-2020 Matt Dunwoodie <ncon@noconroy.net>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#ifndef __NOISE_H__
#define __NOISE_H__

#include <sys/types.h>
#include <sys/time.h>
#include <sys/rwlock.h>

#include <crypto/blake2s.h>
#include <crypto/chachapoly.h>
#include <crypto/curve25519.h>

#define NOISE_PUBLIC_KEY_LEN	CURVE25519_KEY_SIZE
#define NOISE_SYMMETRIC_KEY_LEN	CHACHA20POLY1305_KEY_SIZE
#define NOISE_TIMESTAMP_LEN	(sizeof(uint64_t) + sizeof(uint32_t))
#define NOISE_AUTHTAG_LEN	CHACHA20POLY1305_AUTHTAG_SIZE
#define NOISE_HASH_LEN		BLAKE2S_HASH_SIZE

/* Protocol string constants */
#define NOISE_HANDSHAKE_NAME	"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
#define NOISE_IDENTIFIER_NAME	"WireGuard v1 zx2c4 Jason@zx2c4.com"

/* Constants for the counter */
#define COUNTER_BITS_TOTAL	8192
#define COUNTER_BITS		(sizeof(unsigned long) * 8)
#define COUNTER_NUM		(COUNTER_BITS_TOTAL / COUNTER_BITS)
#define COUNTER_WINDOW_SIZE	(COUNTER_BITS_TOTAL - COUNTER_BITS)

/* Constants for the keypair */
#define REKEY_AFTER_MESSAGES	(1ull << 60)
#define REJECT_AFTER_MESSAGES	(UINT64_MAX - COUNTER_WINDOW_SIZE - 1)
#define REKEY_AFTER_TIME	120
#define REKEY_AFTER_TIME_RECV	165
#define REJECT_AFTER_TIME	180
#define REJECT_INTERVAL		(1000000000 / 50) /* fifty times per sec */
/* 24 = floor(log2(REJECT_INTERVAL)) */
#define REJECT_INTERVAL_MASK	(~((1ull<<24)-1))

enum noise_state_hs {
	HS_ZEROED = 0,
	CREATED_INITIATION,
	CONSUMED_INITIATION,
	CREATED_RESPONSE,
	CONSUMED_RESPONSE,
};

struct noise_handshake {
	enum noise_state_hs	 hs_state;
	uint32_t		 hs_local_index;
	uint32_t		 hs_remote_index;
	uint8_t		 	 hs_e[NOISE_PUBLIC_KEY_LEN];
	uint8_t		 	 hs_hash[NOISE_HASH_LEN];
	uint8_t		 	 hs_ck[NOISE_HASH_LEN];
};

struct noise_counter {
	struct rwlock		 c_lock;
	uint64_t		 c_send;
	uint64_t		 c_recv;
	unsigned long		 c_backtrack[COUNTER_NUM];
};

struct noise_keypair {
	SLIST_ENTRY(noise_keypair)	kp_entry;
	int				kp_valid;
	int				kp_is_initiator;
	uint32_t			kp_local_index;
	uint32_t			kp_remote_index;
	uint8_t				kp_send[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t				kp_recv[NOISE_SYMMETRIC_KEY_LEN];
	struct timespec			kp_birthdate; /* nanouptime */
	struct noise_counter		kp_ctr;
};

struct noise_remote {
	uint8_t				 r_public[NOISE_PUBLIC_KEY_LEN];
	struct noise_local		*r_local;
	uint8_t		 		 r_ss[NOISE_PUBLIC_KEY_LEN];

	struct rwlock			 r_handshake_lock;
	struct noise_handshake		 r_handshake;
	uint8_t				 r_psk[NOISE_SYMMETRIC_KEY_LEN];
	uint8_t				 r_timestamp[NOISE_TIMESTAMP_LEN];
	struct timespec			 r_last_init; /* nanouptime */

	struct rwlock			 r_keypair_lock;
	SLIST_HEAD(,noise_keypair)	 r_unused_keypairs;
	struct noise_keypair		*r_next, *r_current, *r_previous;
	struct noise_keypair		 r_keypair[3]; /* 3: next, current, previous. */

};

struct noise_local {
	struct rwlock		l_identity_lock;
	int			l_has_identity;
	uint8_t			l_public[NOISE_PUBLIC_KEY_LEN];
	uint8_t			l_private[NOISE_PUBLIC_KEY_LEN];

	struct noise_upcall {
		void	 *u_arg;
		struct noise_remote *
			(*u_remote_get)(void *, uint8_t[NOISE_PUBLIC_KEY_LEN]);
		uint32_t
			(*u_index_set)(void *, struct noise_remote *);
		void	(*u_index_drop)(void *, uint32_t);
	}			l_upcall;
};

/* Set/Get noise parameters */
void	noise_local_init(struct noise_local *, struct noise_upcall *);
void	noise_local_lock_identity(struct noise_local *);
void	noise_local_unlock_identity(struct noise_local *);
int	noise_local_set_private(struct noise_local *, uint8_t[NOISE_PUBLIC_KEY_LEN]);
int	noise_local_keys(struct noise_local *, uint8_t[NOISE_PUBLIC_KEY_LEN],
	    uint8_t[NOISE_PUBLIC_KEY_LEN]);

void	noise_remote_init(struct noise_remote *, uint8_t[NOISE_PUBLIC_KEY_LEN],
	    struct noise_local *);
int	noise_remote_set_psk(struct noise_remote *, uint8_t[NOISE_SYMMETRIC_KEY_LEN]);
int	noise_remote_keys(struct noise_remote *, uint8_t[NOISE_PUBLIC_KEY_LEN],
	    uint8_t[NOISE_SYMMETRIC_KEY_LEN]);

/* Should be called anytime noise_local_set_private is called */
void	noise_remote_precompute(struct noise_remote *);

/* Cryptographic functions */
int	noise_create_initiation(
	    struct noise_remote *,
	    uint32_t *s_idx,
	    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
	    uint8_t es[NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN],
	    uint8_t ets[NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN]);

int	noise_consume_initiation(
	    struct noise_local *,
	    struct noise_remote **,
	    uint32_t s_idx,
	    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
	    uint8_t es[NOISE_PUBLIC_KEY_LEN + NOISE_AUTHTAG_LEN],
	    uint8_t ets[NOISE_TIMESTAMP_LEN + NOISE_AUTHTAG_LEN]);

int	noise_create_response(
	    struct noise_remote *,
	    uint32_t *s_idx,
	    uint32_t *r_idx,
	    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
	    uint8_t en[0 + NOISE_AUTHTAG_LEN]);

int	noise_consume_response(
	    struct noise_remote *,
	    uint32_t s_idx,
	    uint32_t r_idx,
	    uint8_t ue[NOISE_PUBLIC_KEY_LEN],
	    uint8_t en[0 + NOISE_AUTHTAG_LEN]);

int	noise_remote_begin_session(struct noise_remote *);
void	noise_remote_clear(struct noise_remote *);
void	noise_remote_expire_current(struct noise_remote *);

int	noise_remote_ready(struct noise_remote *);

int	noise_remote_encrypt(
	    struct noise_remote *,
	    uint32_t *r_idx,
	    uint64_t *nonce,
	    uint8_t *buf,
	    size_t buflen);
int	noise_remote_decrypt(
	    struct noise_remote *,
	    uint32_t r_idx,
	    uint64_t nonce,
	    uint8_t *buf,
	    size_t buflen);

#ifdef WGTEST
void	noise_test();
#endif /* WGTEST */

#endif /* __NOISE_H__ */
