/* cryptoapi correctness (CBC/XTS vs reference + roundtrip) and the
 * shared-session XTS tweak race PoC (userspace build of cryptoapi.c).
 *
 * modes:
 *   ./cryptoapi_test cbc      -> print AES-CBC vectors for openssl cross-match
 *   ./cryptoapi_test xts      -> XTS reference cross-check + roundtrips
 *   ./cryptoapi_test race     -> shared-session concurrency race PoC (XTS vs CBC control)
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <stdbool.h>

#include <crypto/cryptoapi/cryptoapi.h>
#include <crypto/rijndael/rijndael.h>

static unsigned long long rs = 0xB5E1C246B5E1C246ULL;
static unsigned int prng(void){ rs ^= rs<<13; rs ^= rs>>7; rs ^= rs<<17; return (unsigned int)(rs>>16); }
static void phex(const unsigned char*b,int n){for(int i=0;i<n;i++)printf("%02x",b[i]);}
static int fails=0, tests=0;

/* ---------------- XTS reference (IEEE1619, dm-crypt IV convention) ------ */
static void
ref_xts(const unsigned char *key /*32*/, const unsigned char *data, int len,
    uint64_t sector, unsigned char *out, int encrypt)
{
	rijndael_ctx r1e, r1d, r2;
	unsigned char tweak[16];
	memset(tweak, 0, sizeof(tweak));
	for (int i=0;i<8;i++) tweak[i] = (sector >> (8*i)) & 0xff; /* LE64 || 0^64 */
	rijndael_set_key(&r2, key+16, 128);
	rijndael_encrypt(&r2, tweak, tweak);                        /* T0 = E_k2(iv||0) */
	rijndael_set_key(&r1e, key, 128);
	rijndael_set_key(&r1d, key, 128);
	for (int off=0; off<len; off+=16) {
		unsigned char blk[16], blk2[16];
		for (int i=0;i<16;i++) blk[i] = data[off+i] ^ tweak[i];
		if (encrypt) rijndael_encrypt(&r1e, blk, blk2);
		else         rijndael_decrypt(&r1d, blk, blk2);
		for (int i=0;i<16;i++) out[off+i] = blk2[i] ^ tweak[i];
		unsigned carry_in=0, carry_out;
		for (int i=0;i<16;i++){
			carry_out = tweak[i] & 0x80;
			tweak[i] = (tweak[i] << 1) | (carry_in ? 1 : 0);
			carry_in = carry_out;
		}
		if (carry_in) tweak[0] ^= 0x87;
	}
}

/* ---------------- mode: cbc -------------------------------------------- */
static void mode_cbc(void){
	for (int i=0;i<24;i++){
		int klen = (i%3==0)?16:(i%3==1?24:32);
		unsigned char key[32], iv[16], data[64], buf[64];
		for(int j=0;j<klen;j++)key[j]=prng()&0xff;
		for(int j=0;j<16;j++)iv[j]=prng()&0xff;
		for(int j=0;j<64;j++)data[j]=prng()&0xff;
		cryptoapi_cipher_t c = cryptoapi_cipher_find("aes-cbc", klen*8);
		if(!c){printf("FAIL find aes-cbc %d\n",klen);fails++;continue;}
		cryptoapi_cipher_session_t s = cryptoapi_cipher_newsession(c);
		if (cryptoapi_cipher_setkey(s, key, klen)) { printf("FAIL setkey\n"); fails++; continue; }
		memcpy(buf, data, 64);
		if (cryptoapi_cipher_encrypt(s, buf, 64, iv, 16)) { printf("FAIL enc\n"); fails++; }
		printf("CBCVEC %d ", klen); phex(key,klen); printf(" "); phex(iv,16); printf(" "); phex(data,64); printf(" "); phex(buf,64); printf("\n");
		tests++;
		if (cryptoapi_cipher_decrypt(s, buf, 64, iv, 16)) { printf("FAIL dec\n"); fails++; }
		if (memcmp(buf, data, 64)) { printf("FAIL cbc roundtrip\n"); fails++; }
		cryptoapi_cipher_freesession(s);
	}
	printf("SUMMARY cryptoapi_cbc tests=%d fails=%d\n", tests, fails);
}

/* ---------------- mode: xts -------------------------------------------- */
static void mode_xts(void){
	unsigned char key[32];
	for(int j=0;j<32;j++)key[j]=prng()&0xff;
	cryptoapi_cipher_t c = cryptoapi_cipher_find("aes-xts", 256);
	if(!c){printf("FAIL find aes-xts\n");exit(1);}
	cryptoapi_cipher_session_t s = cryptoapi_cipher_newsession(c);
	if (cryptoapi_cipher_setkey(s, key, 32)) { printf("FAIL xts setkey\n"); exit(1); }

	uint64_t sectors[] = {0, 1, 2, 3, 0x1234, 0xffffffffffffffffULL, 0x8000000000000000ULL};
	unsigned char data[512], out[512], ref[512], back[512];
	for (size_t si=0; si<sizeof(sectors)/sizeof(sectors[0]); si++){
		for(int j=0;j<512;j++)data[j]=prng()&0xff;
		unsigned char iv[16]; memset(iv,0,sizeof(iv));
		for(int i=0;i<8;i++) iv[i]=(sectors[si]>>(8*i))&0xff;   /* dm ivgen plain64 */
		memcpy(out,data,512);
		if (cryptoapi_cipher_encrypt(s, out, 512, iv, 16)){printf("FAIL xts enc\n");fails++;}
		ref_xts(key, data, 512, sectors[si], ref, 1);
		tests++;
		if (memcmp(out, ref, 512)) { printf("FAIL xts mismatch sector %llu\n",(unsigned long long)sectors[si]); fails++; }
		memcpy(back,out,512);
		if (cryptoapi_cipher_decrypt(s, back, 512, iv, 16)){printf("FAIL xts dec\n");fails++;}
		if (memcmp(back, data, 512)) { printf("FAIL xts roundtrip sector %llu\n",(unsigned long long)sectors[si]); fails++; }
	}
	/* twofish + serpent roundtrips via cryptoapi */
	const char *names[] = {"twofish-xts","serpent-xts","twofish-cbc","serpent-cbc"};
	int bits[] = {256,256,128,128};
	for (size_t n=0;n<4;n++){
		cryptoapi_cipher_t cc = cryptoapi_cipher_find(names[n], bits[n]);
		if(!cc){printf("FAIL find %s\n",names[n]);fails++;continue;}
		cryptoapi_cipher_session_t ss = cryptoapi_cipher_newsession(cc);
		unsigned char k[64], iv[16];
		for(int j=0;j<bits[n]/8;j++)k[j]=prng()&0xff;
		for(int j=0;j<16;j++)iv[j]=prng()&0xff;
		if (cryptoapi_cipher_setkey(ss, k, bits[n]/8)){printf("FAIL setkey %s\n",names[n]);fails++;continue;}
		unsigned char d2[512], o2[512];
		for(int j=0;j<512;j++)d2[j]=prng()&0xff;
		memcpy(o2,d2,512);
		if (cryptoapi_cipher_encrypt(ss, o2, 512, iv, 16)){printf("FAIL enc %s\n",names[n]);fails++;}
		if (cryptoapi_cipher_decrypt(ss, o2, 512, iv, 16)){printf("FAIL dec %s\n",names[n]);fails++;}
		tests++;
		if (memcmp(o2, d2, 512)) { printf("FAIL roundtrip %s\n", names[n]); fails++; }
		cryptoapi_cipher_freesession(ss);
	}
	cryptoapi_cipher_freesession(s);
	printf("SUMMARY cryptoapi_xts tests=%d fails=%d\n", tests, fails);
}

/* ---------------- mode: race (the finding) ------------------------------ */
#define SECBYTES 512
static unsigned char rkey[32];
static unsigned char secdata[SECBYTES];
static volatile int stop_flag;
static unsigned long long bad_cnt[8];
static unsigned long long iters[8];

/* per-mode reference outputs [2] (sector/iv A and B) */
static unsigned char ref_out[2][SECBYTES];

static void iv_of(uint64_t sector, unsigned char *iv){
	memset(iv,0,16);
	for(int i=0;i<8;i++) iv[i]=(sector>>(8*i))&0xff;
}

struct targ { int idx; int is_xts; cryptoapi_cipher_session_t sess; };
static void *worker(void *v){
	struct targ *t = v;
	unsigned char buf[SECBYTES], iv[16];
	while (!stop_flag){
		memcpy(buf, secdata, SECBYTES);
		iv_of(t->idx & 1, iv);            /* A or B (or fixed IV pair for CBC) */
		cryptoapi_cipher_encrypt(t->sess, buf, SECBYTES, iv, 16);
		iters[t->idx]++;
		if (memcmp(buf, ref_out[t->idx & 1], SECBYTES))
			bad_cnt[t->idx]++;
	}
	return NULL;
}

static unsigned long long run_race(int is_xts, int seconds){
	const char *nm = is_xts ? "aes-xts" : "aes-cbc";
	cryptoapi_cipher_t c = cryptoapi_cipher_find(nm, 256);
	cryptoapi_cipher_session_t s = cryptoapi_cipher_newsession(c);
	if (cryptoapi_cipher_setkey(s, rkey, 32)) exit(1);

	/* references on a separate (uncontended) session */
	for (int i=0;i<2;i++){
		cryptoapi_cipher_session_t r = cryptoapi_cipher_newsession(c);
		cryptoapi_cipher_setkey(r, rkey, 32);
		unsigned char iv[16];
		iv_of(i, iv);
		memcpy(ref_out[i], secdata, SECBYTES);
		cryptoapi_cipher_encrypt(r, ref_out[i], SECBYTES, iv, 16);
		cryptoapi_cipher_freesession(r);
	}
	/* serial sanity on the shared session */
	for (int i=0;i<2;i++){
		unsigned char buf[SECBYTES], iv[16];
		iv_of(i, iv);
		memcpy(buf, secdata, SECBYTES);
		cryptoapi_cipher_encrypt(s, buf, SECBYTES, iv, 16);
		if (memcmp(buf, ref_out[i], SECBYTES)){ printf("SANITY FAIL serial %s %d\n", nm, i); exit(1); }
	}

	memset(bad_cnt,0,sizeof(bad_cnt)); memset(iters,0,sizeof(iters)); stop_flag=0;
	pthread_t th[4]; struct targ ta[4];
	for (int i=0;i<4;i++){ ta[i].idx=i; ta[i].is_xts=is_xts; ta[i].sess=s; pthread_create(&th[i], NULL, worker, &ta[i]); }
	usleep(seconds*1000000);
	stop_flag = 1;
	for (int i=0;i<4;i++) pthread_join(th[i], NULL);
	unsigned long long tot=0, bad=0;
	for (int i=0;i<4;i++){ tot+=iters[i]; bad+=bad_cnt[i]; }
	cryptoapi_cipher_freesession(s);
	printf("RACE %-8s: %llu iterations, %llu corrupted outputs -> %s\n",
	    nm, tot, bad, bad? "CORRUPTION OBSERVED" : "no corruption");
	return bad;
}

static void mode_race(void){
	for(int j=0;j<32;j++)rkey[j]=prng()&0xff;
	for(int j=0;j<SECBYTES;j++)secdata[j]=prng()&0xff;
	printf("serial sanity OK (both modes)\n");
	unsigned long long xts_bad = run_race(1, 3);
	unsigned long long cbc_bad = run_race(0, 3);
	printf("SUMMARY cryptoapi_race corrupted_xts_outputs=%llu corrupted_cbc_outputs=%llu\n", xts_bad, cbc_bad);
	exit(xts_bad ? 3 : 4);
}

int main(int argc, char **argv){
	if (argc<2){fprintf(stderr,"usage: %s cbc|xts|race\n",argv[0]);return 2;}
	if(!strcmp(argv[1],"cbc"))mode_cbc();
	else if(!strcmp(argv[1],"xts"))mode_xts();
	else if(!strcmp(argv[1],"race"))mode_race();
	else return 2;
	return fails?1:0;
}