Remove unnecessary methods from TlsInput (#170)

pfg666 · kostis · web-flow · commit 69501e0cc0c9 · 2025-04-03T09:31:24.000+03:00
* Remove methods from TlsInput that are now unnecessary given the transition to generics.

---------

Co-authored-by: Kostis Sagonas &lt;kostis@it.uu.se&gt;
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/core/TlsSul.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/core/TlsSul.java
@@ -198,7 +198,7 @@ private void initializeTransportHandler() {
                     try {
                         var firstClientHello = transportHandler.fetchData();
                         receivedClientHello = true;
-                        FirstCachedUdpLayer udpLayer = (FirstCachedUdpLayer) context.getState().getTlsContext().getLayerStack().getLowestLayer();
+                        FirstCachedUdpLayer udpLayer = (FirstCachedUdpLayer) context.getTlsContext().getLayerStack().getLowestLayer();
                         udpLayer.setFirstClientHelo(firstClientHello);
                         udpLayer.isFuzzingClient = true;
                     } catch (SocketTimeoutException e) {
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/DtlsOutputMapper.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/DtlsOutputMapper.java
@@ -31,7 +31,7 @@ public DtlsOutputMapper(MapperConfig mapperConfig, TlsOutputBuilder outputBuilde
 
     @Override
     public TlsOutput receiveOutput(TlsExecutionContext context) {
-        TlsContext tlsContext = context.getState().getTlsContext();
+        TlsContext tlsContext = context.getTlsContext();
         try {
             if (tlsContext.getTransportHandler().isClosed()) {
                 return socketClosed();
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/TlsExecutionContext.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/TlsExecutionContext.java
@@ -126,12 +126,4 @@ public Long getWriteRecordNumberEpoch0() {
     protected TlsStepContext buildStepContext() {
         return new TlsStepContext(stepContexts.size());
     }
-
-    /*
-      public Long incrementWriteRecordNumberEpoch0() {
-        Long old = writeRecordNumberEpoch0;
-        writeRecordNumberEpoch0++;
-        return old;
-      }
-    */
 }
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ChangeCipherSpecInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ChangeCipherSpecInput.java
@@ -18,8 +18,8 @@ public ChangeCipherSpecInput() {
 
     @Override
     public void preSendDtlsUpdate(TlsExecutionContext context) {
-        Encryptor encryptor = getState(context).getTlsContext().getRecordLayer().getEncryptor();
-        int writeEpoch = getState(context).getTlsContext().getWriteEpoch();
+        Encryptor encryptor = context.getTlsContext().getRecordLayer().getEncryptor();
+        int writeEpoch = context.getTlsContext().getWriteEpoch();
         Long writeSeqNumForCurrentEpoch = encryptor.getRecordCipher(writeEpoch).getState().getWriteSequenceNumber();
         context.setWriteRecordNumberEpoch0(writeSeqNumForCurrentEpoch + 1);
     }
@@ -35,9 +35,9 @@ public void postSendDtlsUpdate(TlsExecutionContext context) {
         // TLS-Attacker instantiates non-null ciphers even when the pre-master secret has not been yet negotiated.
         // Here, we replace the ciphers instantiated in such cases by null ciphers.
         // This ensures that encrypted messages are more likely to make sense to the SUT.
-        if (getTlsContext(context).getPreMasterSecret() == null) {
-            makeNullCipherAsMostRecent(getTlsContext(context).getRecordLayer().getEncryptor(), getTlsContext(context));
-            makeNullCipherAsMostRecent(getTlsContext(context).getRecordLayer().getDecryptor(), getTlsContext(context));
+        if (context.getTlsContext().getPreMasterSecret() == null) {
+            makeNullCipherAsMostRecent(context.getTlsContext().getRecordLayer().getEncryptor(), context.getTlsContext());
+            makeNullCipherAsMostRecent(context.getTlsContext().getRecordLayer().getDecryptor(), context.getTlsContext());
         }
     }
 
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloInput.java
@@ -48,25 +48,25 @@ public ClientHelloInput(CipherSuite cipherSuite) {
 
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
-        getConfig(context).setDefaultClientSupportedCipherSuites(Arrays.asList(suite));
+        context.getConfig().setDefaultClientSupportedCipherSuites(Arrays.asList(suite));
         if (suite.name().contains("EC")) {
-            getConfig(context).setAddECPointFormatExtension(true);
-            getConfig(context).setAddEllipticCurveExtension(true);
+            context.getConfig().setAddECPointFormatExtension(true);
+            context.getConfig().setAddEllipticCurveExtension(true);
         } else {
-            getConfig(context).setAddECPointFormatExtension(false);
-            getConfig(context).setAddEllipticCurveExtension(false);
+            context.getConfig().setAddECPointFormatExtension(false);
+            context.getConfig().setAddEllipticCurveExtension(false);
         }
 
-        if(getConfig(context).getHighestProtocolVersion().isDTLS13()) {
+        if (context.getConfig().getHighestProtocolVersion().isDTLS13()) {
             // aka supported_groups
-            getConfig(context).setAddEllipticCurveExtension(true);
+            context.getConfig().setAddEllipticCurveExtension(true);
         }
 
         if (resetDigest) {
-            getTlsContext(context).getDigest().reset();
+            context.getTlsContext().getDigest().reset();
         }
 
-        ClientHelloMessage message = new ClientHelloMessage(getConfig(context));
+        ClientHelloMessage message = new ClientHelloMessage(context.getConfig());
 
         // we exclude the sessionId
         if (!withSessionId) {
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloRenegotiationInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloRenegotiationInput.java
@@ -41,30 +41,30 @@ public boolean isEnabled(TlsExecutionContext context) {
         return switch (enabled) {
             case OWN_EPOCH_CHANGE ->
                 // send epoch is 1 or more
-                getTlsContext(context).getWriteEpoch() > 0;
+                context.getTlsContext().getWriteEpoch() > 0;
             case SERVER_EPOCH_CHANGE ->
                 // receive epoch is 1 or more
-                getTlsContext(context).getReadEpoch() > 0;
+                context.getTlsContext().getReadEpoch() > 0;
             case ONCE ->
-                getTlsExecutionContext(context)
+                context
                         .getTlsStepContextStream()
-                        .noneMatch(s -> this.equals(s.getInput()) && s.getIndex() != getTlsExecutionContext(context).getStepCount() - 1);
+                        .noneMatch(s -> this.equals(s.getInput()) && s.getIndex() != context.getStepCount() - 1);
             default ->
                 true;
         };
     }
 
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
-        getTlsContext(context).getDigest().reset();
+        context.getTlsContext().getDigest().reset();
         if (resetMSeq) {
-            getTlsContext(context).setWriteSequenceNumber(getTlsContext(context).getWriteEpoch(), 0);
+            context.getTlsContext().setWriteSequenceNumber(context.getTlsContext().getWriteEpoch(), 0);
         }
-        getTlsContext(context).setReadSequenceNumber(getTlsContext(context).getReadEpoch(), 0);
+        context.getTlsContext().setReadSequenceNumber(context.getTlsContext().getReadEpoch(), 0);
         if (suite != null) {
-            getConfig(context).setDefaultClientSupportedCipherSuites(suite);
+            context.getConfig().setDefaultClientSupportedCipherSuites(suite);
         }
-        ClientHelloMessage message = new ClientHelloMessage(getConfig(context));
+        ClientHelloMessage message = new ClientHelloMessage(context.getConfig());
         if (!isShort) {
             ModifiableByteArray sbyte = new ModifiableByteArray();
             sbyte.setModification(new ByteArrayExplicitValueModification(
@@ -73,7 +73,7 @@ public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
         }
 
         // mbedtls will only engage in renegotiation if the cookie is empty
-        if (!withCookie && getTlsContext(context).getDtlsCookie() != null) {
+        if (!withCookie && context.getTlsContext().getDtlsCookie() != null) {
             ModifiableByteArray sbyte = new ModifiableByteArray();
             sbyte.setModification(new ByteArrayExplicitValueModification(
                     new byte[] {}));
@@ -89,7 +89,7 @@ public void postReceiveUpdate(TlsOutput output, OutputChecker<TlsOutput> abstrac
         switch (enabled) {
             case ON_SERVER_HELLO:
                 if (!TlsOutputChecker.hasServerHello(output)) {
-                    getTlsExecutionContext(context).disableExecution();
+                    context.disableExecution();
                 }
                 break;
             default:
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloWithSessionIdInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientHelloWithSessionIdInput.java
@@ -42,12 +42,12 @@ private void resetTransportHandler(State state) {
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
         // reset and resume the connection
-        resetTransportHandler(getState(context));
+        resetTransportHandler(context.getState().getState());
         if (suite != null) {
-            getConfig(context).setDefaultClientSupportedCipherSuites(suite);
+            context.getConfig().setDefaultClientSupportedCipherSuites(suite);
         }
-        ClientHelloMessage message = new ClientHelloMessage(getConfig(context));
-        message.setSessionId(getTlsContext(context).getChooser()
+        ClientHelloMessage message = new ClientHelloMessage(context.getConfig());
+        message.setSessionId(context.getTlsContext().getChooser()
                 .getServerSessionId());
 
         return new TlsProtocolMessage(message);
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientKeyExchangeInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ClientKeyExchangeInput.java
@@ -32,7 +32,7 @@ public ClientKeyExchangeInput(KeyExchangeAlgorithm algorithm) {
 
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
-        getTlsContext(context).setPreMasterSecret(null);
+        context.getTlsContext().setPreMasterSecret(null);
         if (algorithm == null) {
             throw new RuntimeException("Algorithm not set");
         }
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/DtlsInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/DtlsInput.java
@@ -3,7 +3,6 @@
 import jakarta.xml.bind.annotation.XmlAttribute;
 import org.apache.commons.lang3.NotImplementedException;
 import se.uu.it.dtlsfuzzer.components.sul.mapper.TlsExecutionContext;
-import se.uu.it.dtlsfuzzer.components.sul.mapper.TlsState;
 
 public abstract class DtlsInput extends TlsInput {
 
@@ -32,20 +31,18 @@ protected DtlsInput(String name) {
 
     @Override
     public final void preSendUpdate(TlsExecutionContext context) {
-        TlsState state = getTlsExecutionContext(context).getState();
-
         // if different epoch than current, set the epoch in TLS context
-        if (epoch != null && epoch != state.getTlsContext().getWriteEpoch()) {
-            state.getTlsContext().setWriteEpoch(epoch);
-            contextEpoch = state.getTlsContext().getWriteEpoch();
+        if (epoch != null && epoch != context.getTlsContext().getWriteEpoch()) {
+            context.getTlsContext().setWriteEpoch(epoch);
+            contextEpoch = context.getTlsContext().getWriteEpoch();
         }
 
         // if epoch > 0, deactivate encryption
         if (!encryptionEnabled) {
             throw new NotImplementedException("Disabling encryption is not currently supported.");
         }
 
-        preSendDtlsUpdate(getTlsExecutionContext(context));
+        preSendDtlsUpdate(context);
     }
 
     public void preSendDtlsUpdate(TlsExecutionContext context) {
@@ -55,15 +52,15 @@ public void preSendDtlsUpdate(TlsExecutionContext context) {
     public final void postSendUpdate(TlsExecutionContext context) {
         // reset epoch number and, if original epoch > 0, reactivate encryption
         if (contextEpoch != null) {
-            getTlsExecutionContext(context).getState().getTlsContext().setWriteEpoch(contextEpoch);
+            context.getTlsContext().setWriteEpoch(contextEpoch);
             contextEpoch = null;
         }
 
         if (!encryptionEnabled) {
             throw new NotImplementedException("Re-enabling encryption is not currently supported.");
         }
 
-        postSendDtlsUpdate(getTlsExecutionContext(context));
+        postSendDtlsUpdate(context);
     }
 
     public void postSendDtlsUpdate(TlsExecutionContext context) {
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/FinishedInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/FinishedInput.java
@@ -24,29 +24,29 @@ public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
         // Uncomment line to print digest, TODO remove this when polishing things up
         // System.out.println(ArrayConverter.bytesToHexString(state.getTlsContext().getDigest().getRawBytes()));
         FinishedMessage message = new FinishedMessage();
-        lastSequenceNumber = getTlsContext(context).getWriteSequenceNumber(getTlsContext(context).getWriteEpoch());
-        // getTlsContext(context).setWriteEpoch(getTlsContext(context).getWriteEpoch() + 1);
-        // getTlsContext(context).setWriteSequenceNumber(getTlsContext(context).getWriteEpoch(), 0L);
+        lastSequenceNumber = context.getTlsContext().getWriteSequenceNumber(context.getTlsContext().getWriteEpoch());
+        // context.getTlsContext().setWriteEpoch(context.getTlsContext().getWriteEpoch() + 1);
+        // context.getTlsContext().setWriteSequenceNumber(context.getTlsContext().getWriteEpoch(), 0L);
         return new TlsProtocolMessage(message);
     }
 
     @Override
     public void postSendDtlsUpdate(TlsExecutionContext context) {
-        if (getTlsContext(context).getConfig().getHighestProtocolVersion().isDTLS13() && !getTlsContext(context).isDtls13ShouldSendFinished()){
+        if (context.getTlsContext().getConfig().getHighestProtocolVersion().isDTLS13() && !context.getTlsContext().isDtls13ShouldSendFinished()){
             // invalid Finished message, we shouldn't care
             return;
         }
-        getTlsContext(context).getDigest().reset();
+        context.getTlsContext().getDigest().reset();
         // we have to make this change for learning to scale
-        getTlsContext(context).setWriteSequenceNumber(getTlsContext(context).getWriteEpoch(), lastSequenceNumber + 1);
+        context.getTlsContext().setWriteSequenceNumber(context.getTlsContext().getWriteEpoch(), lastSequenceNumber + 1);
     }
 
     @Override
     public void postReceiveUpdate(TlsOutput output, OutputChecker<TlsOutput> abstractOutputChecker,
             TlsExecutionContext context) {
         if (resetMSeq) {
             if (TlsOutputChecker.hasChangeCipherSpec(output)) {
-                getTlsContext(context).setWriteSequenceNumber(getTlsContext(context).getWriteEpoch(), 0);
+                context.getTlsContext().setWriteSequenceNumber(context.getTlsContext().getWriteEpoch(), 0);
             }
         }
         super.postReceiveUpdate(output, abstractOutputChecker, context);
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/HelloRequestInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/HelloRequestInput.java
@@ -30,11 +30,11 @@ public HelloRequestInput() {
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
         if (resetSequenceNumber) {
-            origMsgSeqNum = getTlsContext(context).getDtlsFragmentLayer().getWriteHandshakeMessageSequence();
-            getTlsContext(context).getDtlsFragmentLayer().setWriteHandshakeMessageSequence(0);
+            origMsgSeqNum = context.getTlsContext().getDtlsFragmentLayer().getWriteHandshakeMessageSequence();
+            context.getTlsContext().getDtlsFragmentLayer().setWriteHandshakeMessageSequence(0);
         }
         if (retransmittedCHAsRefusal) {
-            clientRandom = getTlsContext(context).getClientRandom();
+            clientRandom = context.getTlsContext().getClientRandom();
         }
         HelloRequestMessage hvr = new HelloRequestMessage();
         return new TlsProtocolMessage(hvr);
@@ -52,10 +52,10 @@ public void postReceiveUpdate(TlsOutput output, OutputChecker<TlsOutput> abstrac
             if (disableOnRefusal) {
                 context.disableExecution();
             } else if (resetSequenceNumber) {
-                getTlsContext(context).getDtlsFragmentLayer().setWriteHandshakeMessageSequence(origMsgSeqNum);
+                context.getTlsContext().getDtlsFragmentLayer().setWriteHandshakeMessageSequence(origMsgSeqNum);
             }
         } else if (disableOnRefusal && retransmittedCHAsRefusal
-                && Arrays.equals(clientRandom, getTlsContext(context).getClientRandom())) {
+                && Arrays.equals(clientRandom, context.getTlsContext().getClientRandom())) {
             context.disableExecution();
         }
         super.postReceiveUpdate(output, abstractOutputChecker, context);
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ServerHelloInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/ServerHelloInput.java
@@ -37,7 +37,7 @@ public ServerHelloInput(CipherSuite cipherSuite) {
 
     @Override
     public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
-        Config config = getConfig(context);
+        Config config = context.getConfig();
 
         config.setDefaultServerSupportedCipherSuites(Arrays.asList(suite));
         config.setDefaultClientSupportedCipherSuites(suite);
@@ -46,7 +46,7 @@ public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {
             config.setAddEllipticCurveExtension(true);
         } else {
             config.setAddECPointFormatExtension(false);
-            getConfig(context).setAddEllipticCurveExtension(false);
+            config.setAddEllipticCurveExtension(false);
         }
         if (suite.isPsk()) {
             config.setAddClientCertificateTypeExtension(false);
@@ -69,11 +69,10 @@ public CipherSuite getCipherSuite() {
     @Override
     public void postReceiveUpdate(TlsOutput output, OutputChecker<TlsOutput> abstractOutputChecker,
             TlsExecutionContext context) {
-        TlsExecutionContext ctx = getTlsExecutionContext(context);
         if (shortHs && context.isExecutionEnabled()) {
             Pair<ProtocolMessage, Record> lastChPair = null;
             int lastChStepIndex = -1;
-            List<Pair<ProtocolMessage, Record>> msgRecPairs = ctx.getReceivedMessagesAndRecords();
+            List<Pair<ProtocolMessage, Record>> msgRecPairs = context.getReceivedMessagesAndRecords();
             for (int i = 0; i < msgRecPairs.size(); i++) {
                 if (msgRecPairs.get(i).getKey() instanceof ClientHelloMessage) {
                     lastChStepIndex = i;
@@ -85,16 +84,16 @@ public void postReceiveUpdate(TlsOutput output, OutputChecker<TlsOutput> abstrac
 
             // We reset the digest and append to it, in reverse order SH, the last CH before
             // the SH and optionally the HR which prompted the CH (if such a HR was sent).
-            getTlsContext(ctx).getDigest().reset();
+            context.getTlsContext().getDigest().reset();
 
-            if (digestHR && ctx.getStepContext(lastChStepIndex).getInput() instanceof HelloRequestInput) {
-                byte[] hrBytes = ctx.getStepContext(lastChStepIndex).getReceivedRecords().get(0).getCleanProtocolMessageBytes().getValue();
-                getTlsContext(context).getDigest().append(hrBytes);
+            if (digestHR && context.getStepContext(lastChStepIndex).getInput() instanceof HelloRequestInput) {
+                byte[] hrBytes = context.getStepContext(lastChStepIndex).getReceivedRecords().get(0).getCleanProtocolMessageBytes().getValue();
+                context.getTlsContext().getDigest().append(hrBytes);
             }
             byte[] chBytes = lastChPair.getRight().getCleanProtocolMessageBytes().getValue();
-            byte[] shBytes = ctx.getStepContext().getSentRecords().get(0).getCleanProtocolMessageBytes().getValue();
-            getTlsContext(context).getDigest().append(chBytes);
-            getTlsContext(context).getDigest().append(shBytes);
+            byte[] shBytes = context.getStepContext().getSentRecords().get(0).getCleanProtocolMessageBytes().getValue();
+            context.getTlsContext().getDigest().append(chBytes);
+            context.getTlsContext().getDigest().append(shBytes);
         }
         super.postReceiveUpdate(output, abstractOutputChecker, context);
     }
diff --git a/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/TlsInput.java b/src/main/java/se/uu/it/dtlsfuzzer/components/sul/mapper/symbols/inputs/TlsInput.java

Original file line number	Diff line number	Diff line change
`@@ -126,12 +126,4 @@ public Long getWriteRecordNumberEpoch0() {`
`126`	`126`	`protected TlsStepContext buildStepContext() {`
`127`	`127`	`return new TlsStepContext(stepContexts.size());`
`128`	`128`	`}`
`129`		`-`
`130`		`- /*`
`131`		`- public Long incrementWriteRecordNumberEpoch0() {`
`132`		`- Long old = writeRecordNumberEpoch0;`
`133`		`- writeRecordNumberEpoch0++;`
`134`		`- return old;`
`135`		`- }`
`136`		`- */`
`137`	`129`	`}`
Original file line number	Diff line number	Diff line change
`@@ -32,7 +32,7 @@ public ClientKeyExchangeInput(KeyExchangeAlgorithm algorithm) {`
`32`	`32`
`33`	`33`	`@Override`
`34`	`34`	`public TlsProtocolMessage generateProtocolMessage(TlsExecutionContext context) {`
`35`		`- getTlsContext(context).setPreMasterSecret(null);`
	`35`	`+ context.getTlsContext().setPreMasterSecret(null);`
`36`	`36`	`if (algorithm == null) {`
`37`	`37`	`throw new RuntimeException("Algorithm not set");`
`38`	`38`	`}`