aboutsummaryrefslogblamecommitdiffstats
path: root/src/wireguard.c
blob: c830d102f3d8b738d43ba795f13fb2e036169446 (plain) (tree)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16















                                                                           
                      
                       
                      
                       
                       
                     
                           
                                
 


                              
 
                          
 
                                   



                                                               

                                                   



                                    
 


                                                                             



                                                  
 





                                                               
 
                                                              
                                                                
 















                                                                                           


                                          
                     
    
                                     
                                                               
 






                                                        
 




                                                                           
 


                              
                                     
                                          
 
                                     
                                      
                                                    
                                     
                                    

                                                                
                    

 






                                                                        
                                        
 
                                     
                                         
                                                                   
                                        
                                    




                       
                                                                   
 

                                                                        

                              
                                    


                                                                          


                                    
         






                                   
 









                                                        


    
                                 
 
                                     


    
                                 
 
                                          

 
    
                                  
 
                                                
                                                        

                                                         





                                                         






                                                  

 
    
                                          
 
                                        


    
                                          
 
                                             

 
    
                                           
 
                                                           
                                                                               

                                                               
                                                  


    
                                                                        
 
                                       







                                                                          
                                      
                               

                                         
                                     


                                             

 

                                              
 
                                       




                                                                             
                                         
                                                         




                                                                              
                

                                                
         

                                   
                                        
 

                                             



                                                            
                                        


                                             

 

                                        
 
                                   
 
                                     
                                     
                                    


                                        
                       

 
                   
                                        
 
                                   





                                                                               
                                    

                                        
                       

 
    







                                                               
                                                               
 


                                      


    
                                                               
 


                                     

 

                                            
 





                                                    


    
                                   
 
                                            

                                      


                                        





                                      







                                        

 



                                                                              
 




                                   
 
                                  
 
                                    
 

                                                            
 


                                                                             
 



                                                                         
 


                                                                  
 

                                                                     
 


                                                             
 


                                                                        
 

                                                                      
 

                                                               
 

                                                                 
 

                                                                   
 
                                   
 
                                                





                                                    
 
                                              
                                                                
 

                                                
                          
 





                                        
 
             

                                                                          
 


                                   
                                  
 

                                                                           
 



                                                         
                                    
 

                                                               
 
                                                            
 

                                                                         
 

                                                             
 
                                                     
 

                                                                                
 
                                                            
 



                                                                        
 

                                                               
 

                                            
                                                  
 


                                        
 
                                    
                     

                   
      


                                        
                   
 

 
             

                                                                        
 


                                      
 
                                  
 

                                                                             
 

                                                 
 

                                                           
                                    
 

                                                                                
 


                                                                                  
 


                                                       
 
                     
      

                                                
                   

 


                                                                           
 







                                                                          

                                    
                                        
 


                                                     
 


                                                                              
 


                                                                              
 


                                                                
 
                                                
                                                              
 

                                                                                
                                                                                       
 



                                       
 
 
 
             

                                                                              
 



                                   
                                  
 

                                                           
 
                                     
 

                                                                              
 

                                                       
 

                                             
 
                                          
 


                                                                
 


                                                                              
 



                                                                               
 






                                                                          
 














                                                                                      


                                    
                                              


                                    


                                                    


                                                
                          

                     
                                    
                          
                   

 
             

                                                                          
 

                                  


                                   
                                                               
                             
 

                                                
 

                                                         
 


                                              
 
                                   
 

                                                                               
 
                                                                
 
                                                          
                                                                       
 
                                                    
 



                                                           
                                                                                              



                                                                      
                                                                                        

                                               
                                                                                   
                                                                               
                                                              
 
                                                  
                     


                                               
      


                                               
                   

 
             

                                                                        
 

                                          
 



                                                                           
                                  
                                   
 
                                                               
                             
 
                                        
 


                                                     
 

                                                                              

                                     

                                             
                                                                
 

                                                                                  

                                               


                                                                              
                                                                                       
 

                                           
                     
      
                                       
                   

 




                                                                           
 

                                                 
 
                                                                               
 

                       
 


                                                                           
 

                       
 


                                                                                          
 

                       
 



                                                                                          
 







                                                                                                 
 
 
    
                                                                  
 

                                                                   
 







                                                                                        
 













                                                                                                
 



                                                                                                

 
    
                                               
 



                                                                        

 










                                                             
 
                
                                     
 
                                   
 




                                                   
                                                                           
                                         
                                                                              
                                       
                                                                            
                                     
                                                                               
                                        
            
                                      
 
 

















                                                                        

    






















                                                                              
 










                                                                              
 















                                                                                                     
















                                                                            










                                   





                                                              





                                                      
    





                                                                         








































                                                                                       
/*
 * Copyright (c) 2019 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.
 */

#include <sys/types.h>
#include <sys/malloc.h>
#include <sys/systm.h>
#include <sys/rwlock.h>
#include <sys/refcnt.h>
#include <sys/time.h>
#include <sys/antireplay.h>
#include <lib/libkern/libkern.h>

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

#include <net/wireguard.h>

/* Constants for the wg protocol */
#define WG_CONSTRUCTION "Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
#define WG_IDENTIFIER "WireGuard v1 zx2c4 Jason@zx2c4.com"
#define WG_MAC1 "mac1----"
#define WG_COOKIE "cookie--"

/* First byte indicating packet type on the wire */
#define WG_MSG_INITIATION htole32(1)
#define WG_MSG_RESPONSE htole32(2)
#define WG_MSG_COOKIE htole32(3)
#define WG_MSG_TRANSPORT htole32(4)

/* Ensure struct sizes are correct at compile time. This allows us to overlay
 * the struct on the incoming network packet and access it as if it was a
 * struct. */
CTASSERT(sizeof(struct wg_msg_initiation) == 148);
CTASSERT(sizeof(struct wg_msg_response) == 92);
CTASSERT(sizeof(struct wg_msg_cookie) == 64);
CTASSERT(sizeof(struct wg_msg_transport) == 16);

CTASSERT(WG_KEY_SIZE == CHACHA20POLY1305_KEY_SIZE);
CTASSERT(WG_HASH_SIZE == BLAKE2S_HASH_SIZE);
CTASSERT(WG_KEY_SIZE == BLAKE2S_KEY_SIZE);
CTASSERT(WG_KEY_SIZE == WG_HASH_SIZE);
CTASSERT(WG_MSG_PADDING_SIZE == CHACHA20POLY1305_AUTHTAG_SIZE);
CTASSERT(WG_XNONCE_SIZE == XCHACHA20POLY1305_NONCE_SIZE);

#define ret_error(err) do { ret = err; goto leave; } while (0)
#define offsetof(type, member)  ((size_t)(&((type *)0)->member))

/* wireguard.c */
void wg_kdf(uint8_t [WG_HASH_SIZE], uint8_t [WG_HASH_SIZE],
    uint8_t [WG_HASH_SIZE], uint8_t [WG_KEY_SIZE], uint8_t *, size_t);
void wg_hash2(uint8_t [WG_HASH_SIZE], uint8_t *, size_t, uint8_t *, size_t);
void wg_mix_hash(struct wg_handshake *, uint8_t *, size_t);
void wg_mix_dh(struct wg_handshake *, uint8_t [WG_KEY_SIZE], uint8_t [WG_KEY_SIZE]);
void wg_mix_psk(struct wg_handshake *, uint8_t [WG_KEY_SIZE]);
void wg_handshake_encrypt(struct wg_handshake *, uint8_t *, uint8_t *, size_t);
int wg_handshake_decrypt(struct wg_handshake *, uint8_t *, uint8_t *, size_t);
void wg_timestamp_get(uint8_t [WG_TIMESTAMP_SIZE]);
int wg_timespec_timedout(struct timespec *, time_t);
enum wg_pkt_type wg_pkt_type(uint8_t *, size_t);
void wg_keypair_generate(struct wg_keypair *);
enum wg_error wg_msg_initiation_valid_mac2(struct wg_msg_initiation *, struct wg_cookie *);
enum wg_error wg_msg_response_valid_mac2(struct wg_msg_response *, struct wg_cookie *);

void wg_session_drop(struct wg_session *);
void wg_session_ref(struct wg_session *);

/* Some crappy API */
void
wg_device_init(struct wg_device *dev,
    void (*outq)(struct wg_peer *, enum wg_pkt_type, uint32_t))
{
	bzero(dev, sizeof(*dev));
	dev->d_outq = outq;
	fm_init(&dev->d_peers);
	fm_init(&dev->d_sessions);
	rw_init(&dev->d_lock, "wg_device");
	/* d_cookie_maker, d_keypair initialised to 0 */
}

struct wg_peer *
wg_device_new_peer(struct wg_device *dev, struct wg_pubkey *key, void *arg)
{
	struct wg_peer *peer;
	peer = malloc(sizeof(*peer), M_DEVBUF, M_WAITOK | M_ZERO);

	peer->p_arg = arg;
	peer->p_device = dev;
	peer->p_remote = *key;
	refcnt_init(&peer->p_refcnt);
	rw_init(&peer->p_lock, "wg_peer");

	rw_enter_write(&dev->d_lock);
	rw_enter_write(&peer->p_lock);
	peer->p_id = fm_insert(&dev->d_peers, peer);
	rw_exit_write(&peer->p_lock);
	rw_exit_write(&dev->d_lock);

	/* All other elements of wg_peer are nulled by M_ZERO */
	return peer;
}

struct wg_session *
wg_device_new_session(struct wg_device *dev)
{
	struct wg_session *session;
	session = malloc(sizeof(*session), M_DEVBUF, M_WAITOK | M_ZERO);

	getnanotime(&session->s_created);
	refcnt_init(&session->s_refcnt);

	rw_enter_write(&dev->d_lock);
	rw_enter_write(&session->s_lock);
	session->s_local_id = fm_insert(&dev->d_sessions, session);
	rw_exit_write(&session->s_lock);
	rw_exit_write(&dev->d_lock);

	return session;
}

struct wg_peer *
wg_device_ref_peerkey(struct wg_device *dev, struct wg_pubkey *key)
{
	/* For the time being, we just iterate through peers to find the
	 * matching peer */
	struct wg_peer *peer;
	struct map_item *item;
	rw_enter_read(&dev->d_lock);
	FM_FOREACH_FILLED(item, &dev->d_peers) {
		peer = item->value;
		if (memcmp(key->k, peer->p_remote.k, sizeof(key->k)) == 0)
			break;
		else
			peer = NULL;
	}
	rw_exit_read(&dev->d_lock);

	if (peer)
		wg_peer_ref(peer);

	return peer;
}

struct wg_peer *
wg_device_ref_peerid(struct wg_device *dev, uint32_t id)
{
	struct wg_peer *peer;
	rw_enter_read(&dev->d_lock);
	peer = fm_lookup(&dev->d_peers, id);
	if (peer)
		refcnt_take(&peer->p_refcnt);
	rw_exit_read(&dev->d_lock);
	return peer;
}

void
wg_peer_ref(struct wg_peer *peer)
{
	refcnt_take(&peer->p_refcnt);
}

void
wg_peer_put(struct wg_peer *peer)
{
	refcnt_rele_wake(&peer->p_refcnt);
}

void
wg_peer_drop(struct wg_peer *peer)
{
	rw_enter_write(&peer->p_device->d_lock);
	fm_remove(&peer->p_device->d_peers, peer->p_id);
	rw_exit_write(&peer->p_device->d_lock);
	refcnt_finalize(&peer->p_refcnt, "wg_peer drop");
	free(peer, M_DEVBUF, sizeof(*peer));
}

struct wg_session *
wg_device_ref_session(struct wg_device *dev, uint32_t id)
{
	struct wg_session *session;
	rw_enter_read(&dev->d_lock);
	session = fm_lookup(&dev->d_sessions, id);
	if (session)
		refcnt_take(&session->s_refcnt);
	rw_exit_read(&dev->d_lock);
	return session;
}

void
wg_session_ref(struct wg_session *session)
{
	refcnt_take(&session->s_refcnt);
}

void
wg_session_put(struct wg_session *session)
{
	refcnt_rele_wake(&session->s_refcnt);
}

void
wg_session_drop(struct wg_session *session)
{
	rw_enter_write(&session->s_peer->p_device->d_lock);
	fm_remove(&session->s_peer->p_device->d_sessions, session->s_local_id);
	rw_exit_write(&session->s_peer->p_device->d_lock);
	refcnt_finalize(&session->s_refcnt, "wg_session drop");
	free(session, M_DEVBUF, sizeof(*session));
}

void
wg_peer_attach_session(struct wg_peer *peer, struct wg_session *session)
{
	struct wg_session *old_session;
	/* Assert the session is not attached to another peer.
	 * Assert the session is for the same device that the peer is for.
	 * Assert the session is newly created, by checking s_state. */
	KASSERT(session->s_peer == NULL);
	//KASSERT(session->s_device == peer->p_device);
	KASSERT(session->s_state == WG_STATE_RECV_INITIATION ||
	    session->s_state == WG_STATE_MADE_INITIATION);

	rw_enter_write(&peer->p_lock);
	session->s_peer = peer;
	old_session = peer->p_hs_session;
	peer->p_hs_session = session;
	rw_exit_write(&peer->p_lock);

	if (old_session)
		wg_session_drop(old_session);
}

void
wg_session_promote(struct wg_session *session)
{
	struct wg_session *old_session;
	struct wg_peer *peer = session->s_peer;
	struct wg_keyset *ks = &session->s_keyset;
	struct wg_handshake *hs = &session->s_handshake;

	/* Setup session: derive keys, initialise the antireplay structure */
	rw_enter_write(&session->s_lock);
	if (session->s_state == WG_STATE_RECV_RESPONSE) {
		session->s_state = WG_STATE_INITIATOR;
		wg_kdf(ks->k_txkey.k, ks->k_rxkey.k, NULL, hs->h_ck, NULL, 0);
	} else if (session->s_state == WG_STATE_MADE_RESPONSE) {
		session->s_state = WG_STATE_RESPONDER;
		wg_kdf(ks->k_rxkey.k, ks->k_txkey.k, NULL, hs->h_ck, NULL, 0);
	} else {
		rw_exit_write(&session->s_lock);
		return;
	}

	antireplay_init(&ks->k_ar);
	rw_exit_write(&session->s_lock);

	rw_enter_write(&session->s_lock);
	old_session = peer->p_ks_session_old;
	if (peer->p_ks_session != NULL)
		peer->p_ks_session_old = peer->p_ks_session;
	peer->p_ks_session = peer->p_hs_session;
	peer->p_hs_session = NULL;
	rw_exit_write(&session->s_lock);

	if (old_session != NULL)
		wg_session_drop(old_session);
}

struct wg_session *
wg_peer_hs_session(struct wg_peer *peer)
{
	struct wg_session *session;

	rw_enter_read(&peer->p_lock);
	session = peer->p_hs_session;
	rw_exit_read(&peer->p_lock);

	if (session)
		wg_session_ref(session);
	return session;
}

struct wg_session *
wg_peer_ks_session(struct wg_peer *peer)
{
	struct wg_session *session;
	rw_enter_read(&peer->p_lock);
	session = peer->p_ks_session;
	if (session != NULL &&
	    (wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) ||
	    session->s_keyset.k_txcounter > WG_REJECT_AFTER_MESSAGES))
		session = NULL;
	rw_exit_read(&peer->p_lock);
	if (session != NULL)
		wg_session_ref(session);
	return session;
}

void
wg_device_setkey(struct wg_device *dev, struct wg_privkey *key)
{
	rw_enter_write(&dev->d_lock);
	wg_keypair_from_key(&dev->d_keypair, key);
	rw_exit_write(&dev->d_lock);
}

void
wg_peer_setshared(struct wg_peer *peer, struct wg_privkey *key)
{
	rw_enter_write(&peer->p_lock);
	peer->p_shared = *key;
	rw_exit_write(&peer->p_lock);
}

void
wg_peer_getshared(struct wg_peer *peer, struct wg_privkey *key)
{
	rw_enter_read(&peer->p_lock);
	*key = peer->p_shared;
	rw_exit_read(&peer->p_lock);
}

struct timespec
wg_peer_last_handshake(struct wg_peer *peer)
{
	struct timespec ret = { 0, 0 };
	rw_enter_read(&peer->p_lock);
	if (peer->p_ks_session)
		ret = peer->p_ks_session->s_created;
	rw_exit_read(&peer->p_lock);
	return ret;
}

void
wg_peer_clean(struct wg_peer *peer)
{
	struct wg_session *hs, *ks, *ks_old;
	rw_enter_write(&peer->p_lock);

	hs = peer->p_hs_session;
	ks = peer->p_ks_session;
	ks_old = peer->p_ks_session_old;

	peer->p_hs_session = NULL;
	peer->p_ks_session = NULL;
	peer->p_ks_session_old = NULL;

	rw_enter_write(&peer->p_lock);

	if (hs != NULL)
		wg_session_drop(hs);
	if (ks != NULL)
		wg_session_drop(ks);
	if (ks_old != NULL)
		wg_session_drop(ks_old);

}

/* Crypto  */
enum wg_error
wg_device_rx_initiation(struct wg_device *dev, struct wg_msg_initiation *init,
    struct wg_session **s)
{
	struct wg_peer *peer;
	struct wg_handshake hs;
	struct wg_timestamp ts;
	struct wg_pubkey remote;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	rw_enter_read(&dev->d_lock);

	memcpy(hs.h_remote.k, init->ephemeral, WG_KEY_SIZE);
	wg_keypair_generate(&hs.h_local);

	wg_hash2(hs.h_ck, WG_CONSTRUCTION, strlen(WG_CONSTRUCTION), NULL, 0);
	memcpy(hs.h_hash, hs.h_ck, WG_HASH_SIZE);
	wg_mix_hash(&hs, WG_IDENTIFIER, strlen(WG_IDENTIFIER));

	wg_mix_hash(&hs, dev->d_keypair.pub.k, WG_KEY_SIZE);
	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_dh(&hs, dev->d_keypair.priv.k, hs.h_remote.k);

	if (!wg_handshake_decrypt(&hs, remote.k, init->static_pub,
	    WG_ENCRYPTED_SIZE(sizeof(remote.k))))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, init->static_pub, sizeof(init->static_pub));
	wg_mix_dh(&hs, dev->d_keypair.priv.k, remote.k);

	if (!wg_handshake_decrypt(&hs, ts.t, init->timestamp,
	    WG_ENCRYPTED_SIZE(sizeof(ts.t))))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, init->timestamp, sizeof(init->timestamp));
	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), dev->d_keypair.pub.k,
	    WG_KEY_SIZE);

	blake2s(hs.h_mac, (void *)init, hs.h_k, sizeof(hs.h_mac),
	    offsetof(struct wg_msg_initiation, mac1), sizeof(hs.h_k));

	if (timingsafe_bcmp(hs.h_mac, init->mac1, WG_MAC_SIZE))
		ret_error(WG_MAC);

	if ((peer = wg_device_ref_peerkey(dev, &remote)) == NULL)
		ret_error(WG_UNKNOWN_PEER);

	if (memcmp(&ts, &peer->p_timestamp, WG_TIMESTAMP_SIZE) < 0)
		ret_error(WG_TIMESTAMP);

	rw_exit_read(&dev->d_lock);

	/* Create new session and add to peer */
	session = wg_device_new_session(dev);
	session->s_handshake = hs;
	session->s_remote_id = init->sender;
	session->s_state = WG_STATE_RECV_INITIATION;
	explicit_bzero(&hs, sizeof(hs));
	*s = session;

	wg_peer_attach_session(peer, session);
	dev->d_outq(peer, WG_PKT_RESPONSE, session->s_local_id);

	/* Take reference to return to caller */
	wg_session_ref(session);
	wg_peer_put(peer);

	return WG_OK;
leave:
	rw_exit_read(&dev->d_lock);
	explicit_bzero(&hs, sizeof(hs));
	return ret;
}

enum wg_error
wg_device_rx_response(struct wg_device *dev, struct wg_msg_response *resp,
    struct wg_session **s)
{
	struct wg_handshake hs;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = wg_device_ref_session(dev, resp->receiver)) == NULL)
		return WG_ID;

	rw_enter_read(&dev->d_lock);
	rw_enter_write(&session->s_lock);

	if (session->s_state != WG_STATE_MADE_INITIATION)
		ret_error(WG_STATE);

	/* Make a local copy so we don't clobber the real HS */
	hs = session->s_handshake;

	memcpy(hs.h_remote.k, resp->ephemeral, WG_KEY_SIZE);

	wg_kdf(hs.h_ck, NULL, NULL, hs.h_ck, hs.h_remote.k, WG_KEY_SIZE);
	wg_mix_hash(&hs, hs.h_remote.k, WG_KEY_SIZE);

	wg_mix_dh(&hs, hs.h_local.priv.k, hs.h_remote.k);
	wg_mix_dh(&hs, dev->d_keypair.priv.k, hs.h_remote.k);

	wg_mix_psk(&hs, session->s_peer->p_shared.k);

	if (!wg_handshake_decrypt(&hs, NULL, resp->empty, WG_ENCRYPTED_SIZE(0)))
		ret_error(WG_DECRYPT);

	wg_mix_hash(&hs, resp->empty, WG_ENCRYPTED_SIZE(0));

	wg_hash2(hs.h_k, WG_MAC1, strlen(WG_MAC1), dev->d_keypair.pub.k,
	    WG_KEY_SIZE);
	blake2s(hs.h_mac, (void *)resp, hs.h_k, sizeof(hs.h_mac),
	    offsetof(struct wg_msg_response, mac1), sizeof(hs.h_k));

	if (timingsafe_bcmp(hs.h_mac, resp->mac1, WG_MAC_SIZE))
		ret_error(WG_MAC);

	session->s_handshake = hs;
	session->s_remote_id = resp->sender;
	session->s_state = WG_STATE_RECV_RESPONSE;

	rw_exit_write(&session->s_lock);
	rw_exit_read(&dev->d_lock);
	explicit_bzero(&hs, sizeof(hs));

	wg_session_promote(session);
	*s = session;

	return ret;
leave:
	rw_exit_write(&session->s_lock);
	rw_exit_read(&dev->d_lock);
	explicit_bzero(&hs, sizeof(hs));
	return ret;

}

enum wg_error
wg_device_rx_cookie(struct wg_device *dev, struct wg_msg_cookie *cookie,
    struct wg_session **s)
{
	uint8_t key[WG_KEY_SIZE];
	uint8_t value[WG_COOKIE_SIZE];
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((session = wg_device_ref_session(dev, cookie->receiver)) == NULL)
		return WG_ID;

	rw_enter_write(&session->s_peer->p_lock);
	rw_enter_read(&session->s_lock);

	if (session->s_state != WG_STATE_MADE_INITIATION &&
	    session->s_state != WG_STATE_MADE_RESPONSE)
		ret_error(WG_STATE);

	wg_hash2(key, WG_COOKIE, strlen(WG_COOKIE), session->s_peer->p_remote.k,
	    WG_KEY_SIZE);

	if(!xchacha20poly1305_decrypt(value, cookie->value, sizeof(cookie->value),
	    session->s_handshake.h_mac, WG_MAC_SIZE, cookie->nonce, key))
		ret_error(WG_DECRYPT);

	memcpy(session->s_peer->p_cookie.cookie, value,
	    sizeof(session->s_peer->p_cookie.cookie));
	getnanotime(&session->s_peer->p_cookie.time);

	*s = session;
leave:
	rw_exit_read(&session->s_lock);
	rw_exit_write(&session->s_peer->p_lock);
	return ret;
}

enum wg_error
wg_device_rx_transport(struct wg_device *dev, struct wg_msg_transport *msg,
    size_t len, struct wg_session **s)
{
	enum wg_error ret = WG_OK;
	size_t data_len = len - offsetof(struct wg_msg_transport, data);
	uint64_t counter = letoh64(msg->counter);
	struct wg_session *session;

	if ((session = wg_device_ref_session(dev, msg->receiver)) == NULL)
		return WG_ID;

	wg_session_promote(session);

	rw_enter_read(&session->s_lock);

	if (session->s_state != WG_STATE_INITIATOR &&
	    session->s_state != WG_STATE_RESPONDER)
		ret_error(WG_STATE);

	if (wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) ||
	    session->s_keyset.k_rxcounter > WG_REJECT_AFTER_MESSAGES)
		ret_error(WG_REJECT);

	if (!chacha20poly1305_decrypt(msg->data, msg->data, data_len, NULL, 0,
	    msg->counter, session->s_keyset.k_rxkey.k))
		ret_error(WG_DECRYPT);

	/* TODO fix read lock -> write lock */
	if (antireplay_update(&session->s_keyset.k_ar, counter))
		ret_error(WG_REPLAY);

	session->s_keyset.k_rxcounter = counter;
	session->s_peer->p_rx_bytes += data_len - WG_MAC_SIZE;

	if (session->s_state == WG_STATE_RESPONDER &&
	    wg_timespec_timedout(&session->s_created, WG_REKEY_AFTER_TIME_RECV))
		dev->d_outq(session->s_peer, WG_PKT_INITIATION, session->s_peer->p_id);

	*s = session;
leave:
	rw_exit_read(&session->s_lock);
	return ret;
}


enum wg_error
wg_device_tx_initiation(struct wg_device *dev, struct wg_msg_initiation *init,
    uint32_t id, struct wg_session **s)
{
	struct wg_peer *peer;
	struct wg_handshake *hs;
	struct wg_session *session;

	enum wg_error ret = WG_OK;

	if ((peer = wg_device_ref_peerid(dev, id)) == NULL)
		return WG_ID;

	rw_enter_read(&peer->p_lock);

	if (!wg_timespec_timedout(&peer->p_last_initiation, WG_REKEY_TIMEOUT))
		ret_error(WG_HS_RATE);

	if (peer->p_attempts >= WG_REKEY_ATTEMPT_COUNT)
		ret_error(WG_HS_ATTEMPTS);

	session = wg_device_new_session(dev);
	hs = &session->s_handshake;

	wg_keypair_generate(&hs->h_local);

	init->type = WG_MSG_INITIATION;
	init->sender = session->s_local_id;
	memcpy(init->ephemeral, hs->h_local.pub.k, WG_KEY_SIZE);

	wg_hash2(hs->h_ck, WG_CONSTRUCTION, strlen(WG_CONSTRUCTION), NULL, 0);
	memcpy(hs->h_hash, hs->h_ck, WG_HASH_SIZE);
	wg_mix_hash(hs, WG_IDENTIFIER, strlen(WG_IDENTIFIER));

	wg_mix_hash(hs, peer->p_remote.k, WG_KEY_SIZE);
	wg_kdf(hs->h_ck, NULL, NULL, hs->h_ck, hs->h_local.pub.k, WG_KEY_SIZE);
	wg_mix_hash(hs, hs->h_local.pub.k, WG_KEY_SIZE);
	wg_mix_dh(hs, hs->h_local.priv.k, peer->p_remote.k);

	wg_handshake_encrypt(hs, init->static_pub, dev->d_keypair.pub.k,
	    WG_KEY_SIZE);

	wg_mix_hash(hs, init->static_pub, WG_ENCRYPTED_SIZE(WG_KEY_SIZE));
	wg_mix_dh(hs, dev->d_keypair.priv.k, peer->p_remote.k);

	wg_timestamp_get(init->timestamp);

	wg_handshake_encrypt(hs, init->timestamp, init->timestamp, WG_TIMESTAMP_SIZE);
	wg_mix_hash(hs, init->timestamp, WG_ENCRYPTED_SIZE(WG_TIMESTAMP_SIZE));
	wg_hash2(hs->h_k, WG_MAC1, strlen(WG_MAC1), peer->p_remote.k, WG_KEY_SIZE);

	blake2s(init->mac1, (void *) init, hs->h_k, sizeof(init->mac1),
	    offsetof(struct wg_msg_initiation, mac1), sizeof(hs->h_k));
	memcpy(hs->h_mac, init->mac1, sizeof(hs->h_mac));

	if (wg_timespec_timedout(&peer->p_cookie.time, WG_COOKIE_VALID_TIME))
		bzero(init->mac2, WG_MAC_SIZE);
	else
		blake2s(init->mac2, (void *)init, peer->p_cookie.cookie,
		    sizeof(init->mac2), offsetof(struct wg_msg_initiation, mac2),
		    sizeof(peer->p_cookie.cookie));

	rw_exit_read(&peer->p_lock);

	/* TODO lock? */
	getnanotime(&peer->p_last_initiation);
	peer->p_attempts++;

	/* Attach session to peer */
	session->s_state = WG_STATE_MADE_INITIATION;
	wg_peer_attach_session(peer, session);
	*s = session;

	/* Take reference to return to caller */
	wg_session_ref(session);
	wg_peer_put(peer);
	return WG_OK;
leave:
	rw_exit_read(&peer->p_lock);
	wg_peer_put(peer);
	return ret;
}

enum wg_error
wg_device_tx_response(struct wg_device *dev, struct wg_msg_response *resp,
    uint32_t id, struct wg_session **s)
{
	enum wg_error ret = WG_OK;

	struct wg_handshake *hs;
	struct wg_session *session;

	if ((session = wg_device_ref_session(dev, id)) == NULL)
		return WG_ID;

	rw_enter_read(&session->s_peer->p_lock);
	rw_enter_write(&session->s_lock);

	if (session->s_state != WG_STATE_RECV_INITIATION)
		ret_error(WG_STATE);

	resp->type = WG_MSG_RESPONSE;
	resp->sender = session->s_local_id;
	resp->receiver = session->s_remote_id;

	hs = &session->s_handshake;

	wg_kdf(hs->h_ck, NULL, NULL, hs->h_ck, hs->h_local.pub.k, WG_KEY_SIZE);
	wg_mix_hash(hs, hs->h_local.pub.k, WG_KEY_SIZE);

	memcpy(resp->ephemeral, hs->h_local.pub.k, WG_KEY_SIZE);

	wg_mix_dh(hs, hs->h_local.priv.k, hs->h_remote.k);
	wg_mix_dh(hs, hs->h_local.priv.k, session->s_peer->p_remote.k);

	wg_mix_psk(hs, session->s_peer->p_shared.k);

	wg_handshake_encrypt(hs, resp->empty, NULL, 0);

	wg_mix_hash(hs, resp->empty, WG_ENCRYPTED_SIZE(0));

	wg_hash2(hs->h_k, WG_MAC1, strlen(WG_MAC1), session->s_peer->p_remote.k, WG_KEY_SIZE);
	blake2s(resp->mac1, (void *)resp, hs->h_k, sizeof(resp->mac1),
	    offsetof(struct wg_msg_response, mac1), sizeof(hs->h_k));
	memcpy(hs->h_mac, resp->mac1, sizeof(hs->h_mac));

	if (wg_timespec_timedout(&session->s_peer->p_cookie.time, WG_COOKIE_VALID_TIME))
		bzero(resp->mac2, WG_MAC_SIZE);
	else
		blake2s(resp->mac2, (void *)resp, session->s_peer->p_cookie.cookie,
		    sizeof(resp->mac2), offsetof(struct wg_msg_response, mac2),
		    sizeof(session->s_peer->p_cookie.cookie));

	session->s_state = WG_STATE_MADE_RESPONSE;
	*s = session;
	rw_exit_write(&session->s_lock);
	rw_exit_read(&session->s_peer->p_lock);
	return ret;
leave:
	rw_exit_write(&session->s_lock);
	rw_exit_read(&session->s_peer->p_lock);
	wg_session_put(session);
	return ret;
}

enum wg_error
wg_device_tx_cookie(struct wg_device *dev, struct wg_msg_cookie *cookie,
    uint32_t id, struct wg_session **s)
{
	panic("should not call this yet");
}

enum wg_error
wg_device_tx_transport(struct wg_device *dev, struct wg_msg_transport *msg,
    size_t len, uint32_t id, struct wg_session **s)
{
	enum wg_error ret = WG_OK;
	struct wg_session *session;

	if ((session = wg_device_ref_session(dev, id)) == NULL)
		return WG_ID;

	rw_enter_read(&session->s_lock);

	if (session->s_state != WG_STATE_INITIATOR &&
	    session->s_state != WG_STATE_RESPONDER)
		ret_error(WG_STATE);

	if (wg_timespec_timedout(&session->s_created, WG_REJECT_AFTER_TIME) ||
	    session->s_keyset.k_txcounter > WG_REJECT_AFTER_MESSAGES)
		ret_error(WG_REJECT);

	msg->type = WG_MSG_TRANSPORT;
	msg->receiver = session->s_remote_id;
	msg->counter = htole64(session->s_keyset.k_txcounter++);

	chacha20poly1305_encrypt(msg->data, msg->data, len, NULL, 0, msg->counter,
	    session->s_keyset.k_txkey.k);

	/* Packet OK, but we do want a rekey */
	if ((session->s_state == WG_STATE_INITIATOR &&
	    wg_timespec_timedout(&session->s_created, WG_REKEY_AFTER_TIME)) ||
	    session->s_keyset.k_txcounter > WG_REKEY_AFTER_MESSAGES)
		dev->d_outq(session->s_peer, WG_PKT_INITIATION, session->s_peer->p_id);

	session->s_peer->p_tx_bytes += len;

	*s = session;
leave:
	rw_exit_read(&session->s_lock);
	return ret;
}

/* WireGuard crypto helper functions */
void
wg_kdf(uint8_t first[WG_HASH_SIZE], uint8_t second[WG_HASH_SIZE],
    uint8_t third[WG_HASH_SIZE], uint8_t key[WG_KEY_SIZE], uint8_t * input,
    size_t input_len)
{
	uint8_t 	buffer[WG_HASH_SIZE + 1];
	uint8_t 	secret[WG_HASH_SIZE];

	blake2s_hmac(secret, input, key, WG_HASH_SIZE, input_len, WG_KEY_SIZE);

	if (!first)
		return;

	buffer[0] = 1;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, 1, WG_KEY_SIZE);
	memcpy(first, buffer, WG_HASH_SIZE);

	if (!second)
		return;

	buffer[WG_HASH_SIZE] = 2;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, WG_HASH_SIZE + 1, WG_KEY_SIZE);
	memcpy(second, buffer, WG_HASH_SIZE);

	if (!third)
		return;

	buffer[WG_HASH_SIZE] = 3;
	blake2s_hmac(buffer, buffer, secret, WG_HASH_SIZE, WG_HASH_SIZE + 1, WG_KEY_SIZE);
	memcpy(third, buffer, WG_HASH_SIZE);
}

void
wg_hash2(uint8_t out[WG_HASH_SIZE], uint8_t * in1, size_t in1_len, uint8_t * in2, size_t in2_len)
{
	struct blake2s_state s;
	blake2s_init(&s, WG_HASH_SIZE);
	blake2s_update(&s, in1, in1_len);
	blake2s_update(&s, in2, in2_len);
	blake2s_final(&s, out, WG_HASH_SIZE);
}

void
wg_mix_hash(struct wg_handshake * hs, uint8_t * in, size_t in_len)
{
	wg_hash2(hs->h_hash, hs->h_hash, WG_HASH_SIZE, in, in_len);
}

void
wg_mix_dh(struct wg_handshake * hs, uint8_t priv[WG_KEY_SIZE], uint8_t pub[WG_KEY_SIZE])
{
	uint8_t 	DH    [WG_KEY_SIZE];
	crypto_scalarmult_curve25519(DH, priv, pub);
	wg_kdf(hs->h_ck, hs->h_k, NULL, hs->h_ck, DH, WG_KEY_SIZE);
	explicit_bzero(DH, WG_KEY_SIZE);
}

void
wg_mix_psk(struct wg_handshake * hs, uint8_t psk[WG_KEY_SIZE])
{
	uint8_t 	t     [WG_HASH_SIZE];
	wg_kdf(hs->h_ck, t, hs->h_k, hs->h_ck, psk, WG_KEY_SIZE);
	wg_mix_hash(hs, t, WG_HASH_SIZE);
	explicit_bzero(t, WG_HASH_SIZE);
}

void
wg_handshake_encrypt(struct wg_handshake *hs, uint8_t *dst, uint8_t *src, size_t srclen)
{
	return chacha20poly1305_encrypt(dst, src, srclen, hs->h_hash, WG_HASH_SIZE, 0, hs->h_k);
}

int
wg_handshake_decrypt(struct wg_handshake *hs, uint8_t *dst, uint8_t *src, size_t srclen)
{
	return chacha20poly1305_decrypt(dst, src, srclen, hs->h_hash, WG_HASH_SIZE, 0, hs->h_k);
}

void
wg_timestamp_get(uint8_t ts[WG_TIMESTAMP_SIZE])
{
	struct timespec tv;
	getnanotime(&tv);
	*(uint64_t *) (ts) = htobe64(0x400000000000000aULL + tv.tv_sec);
	*(uint32_t *) (ts + 8) = htobe32(tv.tv_nsec);
}

int
wg_timespec_timedout(struct timespec * start, time_t timeout)
{
	struct timespec now;

	getnanotime(&now);

	return now.tv_sec == start->tv_sec + timeout ?
		now.tv_nsec > start->tv_nsec :
		now.tv_sec > start->tv_sec + timeout;
}

enum wg_pkt_type
wg_pkt_type(uint8_t *buf, size_t len)
{
	struct wg_msg_unknown *msg;

	if (len < sizeof(*msg))
		return WG_PKT_UNKNOWN;
	else
		msg = (struct wg_msg_unknown *)buf;

	if (msg->type == WG_MSG_INITIATION && len == wg_pkt_len[msg->type])
		return WG_PKT_INITIATION;
	else if (msg->type == WG_MSG_RESPONSE && len == wg_pkt_len[msg->type])
		return WG_PKT_RESPONSE;
	else if (msg->type == WG_MSG_COOKIE && len == wg_pkt_len[msg->type])
		return WG_PKT_COOKIE;
	else if (msg->type == WG_MSG_TRANSPORT && len >= wg_pkt_len[msg->type])
		return WG_PKT_TRANSPORT;
	else
		return WG_PKT_UNKNOWN;
}

void
wg_keypair_from_key(struct wg_keypair *kp, const struct wg_privkey *key)
{
	const uint8_t basepoint[WG_KEY_SIZE] = {9};

	/* memmove as key may overlap with kp->priv */
	memmove(kp->priv.k, key->k, WG_KEY_SIZE);

	/* We don't care if the input is not a valid key, we just set
	 * the bits and be done with it. The curve25519 library *SHOULD*
	 * do this, but since this may be returned to the user, we do
	 * it here too. */
	kp->priv.k[0] &= 248;
	kp->priv.k[31] &= 127;
	kp->priv.k[31] |= 64;

	crypto_scalarmult_curve25519(kp->pub.k, kp->priv.k, basepoint);
}

void
wg_keypair_generate(struct wg_keypair *kp)
{
	arc4random_buf(kp->priv.k, sizeof(kp->priv.k));
	wg_keypair_from_key(kp, &kp->priv);
}

enum wg_error
wg_msg_initiation_valid_mac2(struct wg_msg_initiation *m, struct wg_cookie *c)
{
	uint8_t mac[WG_MAC_SIZE];
	blake2s(mac, (uint8_t *)m, c->cookie, WG_MAC_SIZE,
	    offsetof(struct wg_msg_initiation, mac2), WG_COOKIE_SIZE);
	return timingsafe_bcmp(mac, m->mac2, WG_MAC_SIZE) ? WG_MAC : WG_OK;
}

enum wg_error
wg_msg_response_valid_mac2(struct wg_msg_response *m, struct wg_cookie *c)
{
	uint8_t mac[WG_MAC_SIZE];
	blake2s(mac, (uint8_t *)m, c->cookie, WG_MAC_SIZE,
	    offsetof(struct wg_msg_response, mac2), WG_COOKIE_SIZE);
	return timingsafe_bcmp(mac, m->mac2, WG_MAC_SIZE) ? WG_MAC : WG_OK;
}

void
wg_cookie_from_token(struct wg_cookie *c, struct wg_cookie_maker *cm,
    uint8_t *ip, uint8_t ip_len)
{
	if (wg_timespec_timedout(&cm->time, WG_COOKIE_VALID_TIME)) {
		getnanotime(&cm->time);
		arc4random_buf(cm->seed, sizeof(cm->seed));
	}

	blake2s(c->cookie, ip, cm->seed, WG_MAC_SIZE, ip_len, WG_COOKIE_SIZE);
}

void
wg_device_make_cookie(struct wg_device *dev, struct wg_cookie *c,
    uint32_t sender, uint8_t mac[WG_MAC_SIZE], struct wg_msg_cookie *m)
{
	uint8_t key[WG_KEY_SIZE]; // Same as WG_HASH_SIZE

	m->type = WG_MSG_COOKIE;
	m->receiver = sender;
	arc4random_buf(m->nonce, sizeof(m->nonce));

	wg_hash2(key, WG_COOKIE, strlen(WG_COOKIE), dev->d_keypair.pub.k, WG_KEY_SIZE);
	xchacha20poly1305_encrypt(m->value, c->cookie, WG_MAC_SIZE, mac, WG_MAC_SIZE, m->nonce, key);
}

/* Timer Functions */
void
wg_timer_setup(struct wg_timers *t, void *p, void (*keepalive)(void *),
    void (*broken)(void *), void (*reinit)(void *), void (*cleanup)(void *))
{
	assert(p != NULL);
	assert(keepalive != NULL);
	assert(broken != NULL);
	assert(reinit != NULL);
	assert(cleanup != NULL);

	timeout_set_proc(&t->t_ka, keepalive, p);
	timeout_set_proc(&t->t_pka, keepalive, p);
	timeout_set_proc(&t->t_broken, broken, p);
	timeout_set_proc(&t->t_reinit, reinit, p);
	timeout_set_proc(&t->t_cleanup, cleanup, p);
}

void
wg_timer_stop(struct wg_timers *t)
{
	/* TODO need barrier? */
	timeout_del(&t->t_ka);
	timeout_del(&t->t_pka);
	timeout_del(&t->t_broken);
	timeout_del(&t->t_reinit);
	timeout_del(&t->t_cleanup);
}

void
wg_timer_persistent_keepalive_tick(struct wg_timers *t)
{
	if (t->t_pka_interval > 0)
		timeout_add_sec(&t->t_pka, t->t_pka_interval);
}

uint16_t
wg_timer_persistent_keepalive_get(struct wg_timers *t)
{
	return t->t_pka_interval;
}

void
wg_timer_persistent_keepalive_set(struct wg_timers *t, uint16_t interval)
{
	t->t_pka_interval = interval;
}

void
wg_timer_cleanup_tick(struct wg_timers *t)
{
	timeout_add_sec(&t->t_cleanup, WG_REJECT_AFTER_TIME * 3);
}

void
wg_timer_keepalive_flag(struct wg_timers *t)
{
	timeout_add_sec(&t->t_ka, WG_KEEPALIVE_TIMEOUT);
}

void
wg_timer_keepalive_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_ka);
}

void
wg_timer_broken_flag(struct wg_timers *t)
{
	if (timeout_pending(&t->t_broken) == 0)
		timeout_add_sec(&t->t_broken, WG_REKEY_TIMEOUT + WG_KEEPALIVE_TIMEOUT);
}

void
wg_timer_broken_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_broken);
}

void
wg_timer_reinit_flag(struct wg_timers *t)
{
	timeout_add_sec(&t->t_reinit, WG_REKEY_TIMEOUT);
}

void
wg_timer_reinit_unflag(struct wg_timers *t)
{
	timeout_del(&t->t_reinit);
}