Skip to content

Commit e4ed9e3

Browse files
committed
RatchetTest: initial mods to test PQ, WIP
1 parent b787e15 commit e4ed9e3

File tree

1 file changed

+99
-24
lines changed

1 file changed

+99
-24
lines changed

java-utils/RatchetTest.java

Lines changed: 99 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import java.util.Arrays;
1717
import java.util.ArrayList;
1818
import java.util.Date;
19+
import java.util.EnumSet;
1920
import java.util.List;
2021

2122
// both the following are GeneralSecurityExceptions
@@ -125,6 +126,7 @@ public class RatchetTest implements RatchetPayload.PayloadCallback {
125126
private final File _riFile;
126127
private final File _skFile;
127128
private final File _pkFile;
129+
private final EncType _type;
128130

129131
private final X25519KeyFactory kf;
130132
private final Elg2KeyFactory ekf;
@@ -135,6 +137,11 @@ public class RatchetTest implements RatchetPayload.PayloadCallback {
135137

136138
public RatchetTest(RouterContext ctx, String host, int port, boolean isAlice, boolean useNoise,
137139
File rifile, File skfile, File pkfile) {
140+
this(ctx, host, port, isAlice, useNoise, rifile, skfile, pkfile, EncType.ECIES_X25519);
141+
}
142+
143+
public RatchetTest(RouterContext ctx, String host, int port, boolean isAlice, boolean useNoise,
144+
File rifile, File skfile, File pkfile, EncType type) {
138145
_context = ctx;
139146
_log = ctx.logManager().getLog(getClass());
140147
_log.setMinimumPriority(Log.DEBUG);
@@ -147,8 +154,14 @@ public RatchetTest(RouterContext ctx, String host, int port, boolean isAlice, bo
147154
_riFile = rifile;
148155
_skFile = skfile;
149156
_pkFile = pkfile;
150-
151-
kf = new X25519KeyFactory(_context);
157+
_type = type;
158+
159+
if (type == EncType.ECIES_X25519)
160+
kf = new X25519KeyFactory(_context);
161+
else if (useNoise)
162+
kf = new TypeKeyFactory(_context, type);
163+
else
164+
throw new IllegalArgumentException("PQ types require -n");
152165
kf.start();
153166
ekf = new Elg2KeyFactory(_context);
154167
ekf.start();
@@ -157,8 +170,42 @@ public RatchetTest(RouterContext ctx, String host, int port, boolean isAlice, bo
157170

158171
if (useNoise)
159172
System.out.println("Using Noise state machine mode");
173+
System.out.println("Using encryption type " + type);
174+
}
175+
176+
/**
177+
* Make type 5-7 keys since we don't have a key factory yet
178+
*/
179+
private static class TypeKeyFactory extends X25519KeyFactory {
180+
private final EncType etype;
181+
public TypeKeyFactory(RouterContext ctx, EncType type) {
182+
super(ctx);
183+
etype = type;
184+
}
185+
@Override
186+
public KeyPair getKeys() {
187+
KeyPair rv = super.getKeys();
188+
PublicKey pk = new PublicKey(etype, rv.getPublic().getData());
189+
PrivateKey sk = new PrivateKey(etype, rv.getPrivate().getData());
190+
rv = new KeyPair(pk, sk);
191+
return rv;
192+
}
160193
}
161194

195+
private static String getNoisePattern(EncType type) {
196+
switch(type) {
197+
case ECIES_X25519:
198+
return HandshakeState.PATTERN_ID_IK;
199+
case MLKEM512_X25519:
200+
return HandshakeState.PATTERN_ID_IKHFS_512;
201+
case MLKEM768_X25519:
202+
return HandshakeState.PATTERN_ID_IKHFS_768;
203+
case MLKEM1024_X25519:
204+
return HandshakeState.PATTERN_ID_IKHFS_1024;
205+
default:
206+
throw new IllegalArgumentException("No pattern for " + type);
207+
}
208+
}
162209

163210
///////////////////////////////////////////////////////
164211
// Alice
@@ -180,17 +227,17 @@ private void connect() throws Exception {
180227
if (keylist == null)
181228
throw new IllegalArgumentException("no keys");
182229
for (PublicKey k : keylist) {
183-
if (k.getType().equals(EncType.ECIES_X25519)) {
230+
if (k.getType().equals(_type)) {
184231
pk = k;
185232
break;
186233
}
187234
}
188235
if (pk == null)
189-
throw new IOException("no ECIES key");
236+
throw new IOException("no " + _type + " key found in Bob's leaseset");
190237
s = pk.getData();
191238
} else {
192239
s = Base64.decode(BOB);
193-
pk = new PublicKey(EncType.ECIES_X25519, s);
240+
pk = new PublicKey(_type, s);
194241
bobhash = Hash.FAKE_HASH;
195242
}
196243
System.out.println("Loaded Bob LS2:\n" /* + bob */);
@@ -205,7 +252,7 @@ private void connect() throws Exception {
205252
byte[] epriv = null;
206253
HandshakeState state = null;
207254
if (_useNoise) {
208-
state = new HandshakeState(HandshakeState.PATTERN_ID_IK, HandshakeState.INITIATOR, ekf);
255+
state = new HandshakeState(getNoisePattern(_type), HandshakeState.INITIATOR, ekf);
209256
state.getRemotePublicKey().setPublicKey(s, 0);
210257
state.getLocalKeyPair().setKeys(priv, 0, as, 0);
211258
System.out.println("Before start");
@@ -242,15 +289,15 @@ private void connect() throws Exception {
242289

243290
// create payload
244291
int padlen = 13;
245-
byte[] payload = new byte[3 + 36 + 3 + padlen];
292+
byte[] payload = new byte[3 + 1 + 3 + padlen];
246293
List<Block> blocks = new ArrayList<Block>(4);
247294
Block block = new AckRequestBlock();
248295
blocks.add(block);
249296
block = new PaddingBlock(padlen);
250297
blocks.add(block);
251298
int payloadlen = createPayload(payload, 0, blocks);
252299
if (payloadlen != payload.length)
253-
throw new IllegalStateException("payload size mismatch");
300+
throw new IllegalStateException("payload size mismatch expected " + payload.length + " got " + payloadlen);
254301
System.out.println("raw payload is:");
255302
System.out.println(HexDump.dump(payload));
256303

@@ -604,13 +651,13 @@ private void listen() throws Exception {
604651
throw new IllegalArgumentException("no keys");
605652
PublicKey pk = null;
606653
for (PublicKey k : keylist) {
607-
if (k.getType().equals(EncType.ECIES_X25519)) {
654+
if (k.getType().equals(_type)) {
608655
pk = k;
609656
break;
610657
}
611658
}
612659
if (pk == null)
613-
throw new IOException("no ECIES key");
660+
throw new IOException("no " + _type + " key found in Bob's leaseset");
614661
s = pk.getData();
615662

616663
KeyPair keys = keyload(_skFile);
@@ -637,7 +684,7 @@ private void listen() throws Exception {
637684

638685

639686
ServerSocket server = new ServerSocket(_port, 0, InetAddress.getByName(_host));
640-
System.out.println("Bob listening for NTCP2 connections on " + _host + ':' + _port);
687+
System.out.println("Bob listening for Ratchet connections on " + _host + ':' + _port);
641688

642689
while (_running) {
643690
Socket socket = null;
@@ -711,7 +758,7 @@ private void runConnection(Socket socket, byte[] s, byte[] inithash, byte[] priv
711758

712759
HandshakeState state = null;
713760
if (_useNoise) {
714-
state = new HandshakeState(HandshakeState.PATTERN_ID_IK, HandshakeState.RESPONDER, ekf);
761+
state = new HandshakeState(getNoisePattern(_type), HandshakeState.RESPONDER, ekf);
715762
state.getLocalKeyPair().setKeys(priv, 0, s, 0);
716763
System.out.println("Before start");
717764
System.out.println(state.toString());
@@ -761,7 +808,7 @@ private void runConnection(Socket socket, byte[] s, byte[] inithash, byte[] priv
761808
// save for tagset HKDF
762809
ck = state.getChainingKey();
763810
state.getRemotePublicKey().getPublicKey(alicestatic, 0);
764-
alicePubkey = new PublicKey(EncType.ECIES_X25519, alicestatic);
811+
alicePubkey = new PublicKey(_type, alicestatic);
765812
} else {
766813
// MixHash(e.pubkey)
767814
System.out.println("MixHash(bob static) " + Base64.encode(inithash));
@@ -789,7 +836,7 @@ private void runConnection(Socket socket, byte[] s, byte[] inithash, byte[] priv
789836
throw new IOException("Bad AEAD msg 1 pt 1", gse);
790837
}
791838
System.out.println("Got Alice static key: " + Base64.encode(alicestatic));
792-
alicePubkey = new PublicKey(EncType.ECIES_X25519, alicestatic);
839+
alicePubkey = new PublicKey(_type, alicestatic);
793840
// MixHash(cipertext)
794841
tmp2 = new byte[48];
795842
System.arraycopy(tmp, 32, tmp2, 0, 48);
@@ -833,15 +880,15 @@ private void runConnection(Socket socket, byte[] s, byte[] inithash, byte[] priv
833880

834881
// create payload
835882
int padlen = 13;
836-
payload = new byte[3 + 36 + 3 + padlen];
883+
payload = new byte[3 + 1 + 3 + padlen];
837884
List<Block> blocks = new ArrayList<Block>(4);
838885
Block block = new AckRequestBlock();
839886
blocks.add(block);
840887
block = new PaddingBlock(padlen);
841888
blocks.add(block);
842889
payloadlen = createPayload(payload, 0, blocks);
843890
if (payloadlen != payload.length)
844-
throw new IllegalStateException("payload size mismatch");
891+
throw new IllegalStateException("payload size mismatch expected " + payload.length + " got " + payloadlen);
845892
System.out.println("payload is:");
846893
System.out.println(HexDump.dump(payload));
847894

@@ -1353,7 +1400,7 @@ public void mixHash(byte[] hash, byte data[], byte[] ohash) {
13531400

13541401

13551402
/** secret keys */
1356-
private static KeyPair keyload(File f) throws Exception {
1403+
private KeyPair keyload(File f) throws Exception {
13571404
byte[] pub = new byte[32];
13581405
byte[] priv = new byte[32];
13591406
InputStream is = null;
@@ -1369,8 +1416,8 @@ private static KeyPair keyload(File f) throws Exception {
13691416
try { is.close(); } catch (IOException ioe) {}
13701417
}
13711418
}
1372-
PublicKey kp = new PublicKey(EncType.ECIES_X25519, pub);
1373-
PrivateKey ks = new PrivateKey(EncType.ECIES_X25519, priv);
1419+
PublicKey kp = new PublicKey(_type, pub);
1420+
PrivateKey ks = new PrivateKey(_type, priv);
13741421
return new KeyPair(kp, ks);
13751422
}
13761423

@@ -1437,9 +1484,10 @@ public static void main(String args[]) throws Exception {
14371484
File sk = null;
14381485
File pk = null;
14391486
File ri = null;
1487+
String type = null;
14401488
boolean useNoise = false;
14411489

1442-
Getopt g = new Getopt("ntcp2test", args, "h:p:s:k:r:n");
1490+
Getopt g = new Getopt("ratchettest", args, "h:p:s:k:r:nt:");
14431491
int c;
14441492
while ((c = g.getopt()) != -1) {
14451493
switch (c) {
@@ -1467,6 +1515,10 @@ public static void main(String args[]) throws Exception {
14671515
useNoise = true;
14681516
break;
14691517

1518+
case 't':
1519+
type = g.getOptarg();
1520+
break;
1521+
14701522
default:
14711523
error = true;
14721524
}
@@ -1483,7 +1535,16 @@ public static void main(String args[]) throws Exception {
14831535
RouterContext ctx = RouterContext.listContexts().get(0);
14841536
boolean isAlice = pk == null && sk == null;
14851537

1486-
RatchetTest test = new RatchetTest(ctx, host, port, isAlice, useNoise, ri, pk, sk);
1538+
EncType etype = EncType.ECIES_X25519;
1539+
if (type != null) {
1540+
etype = EncType.parseEncType(type);
1541+
if (etype == null) {
1542+
System.err.println("Bad enc type " + type);
1543+
System.exit(1);
1544+
}
1545+
}
1546+
1547+
RatchetTest test = new RatchetTest(ctx, host, port, isAlice, useNoise, ri, pk, sk, etype);
14871548
System.out.println("Starting RatchetTest");
14881549
try {
14891550
if (isAlice)
@@ -1499,9 +1560,23 @@ public static void main(String args[]) throws Exception {
14991560

15001561
private static final void usage() {
15011562
System.err.println("Usage: RatchetTest\n" +
1502-
" bob listen args: [-h bobhost] [-p bobport] [-n] -s bobsecretkeyfile -k bobprivkeyfile -r bobls2file\n" +
1503-
" alice connect args: [-h bobhost] [-p bobport] [-n] -r bobls2file\n" +
1563+
" bob listen args: [-h bobhost] [-p bobport] [-t enctype] [-n] -s bobsecretkeyfile -k bobprivkeyfile -r bobls2file\n" +
1564+
" alice connect args: [-h bobhost] [-p bobport] [-t enctype] [-n] -r bobls2file\n" +
15041565
" -n: use noise state machine\n" +
1505-
" default host 127.0.0.1; default port 8887");
1566+
" default host 127.0.0.1; default port 8887\n" +
1567+
" The bob side will create the files if absent.");
1568+
StringBuilder buf = new StringBuilder();
1569+
buf.append("\nAvailable encryption types:\n");
1570+
for (EncType t : EnumSet.allOf(EncType.class)) {
1571+
if (!t.isAvailable())
1572+
continue;
1573+
if (t.getCode() == 0)
1574+
continue;
1575+
buf.append(" ").append(t).append("\t(code: ").append(t.getCode()).append(')');
1576+
if (t.getCode() == 4)
1577+
buf.append(" DEFAULT");
1578+
buf.append('\n');
1579+
}
1580+
System.err.println(buf.toString());
15061581
}
15071582
}

0 commit comments

Comments
 (0)