Repository: xlate/staedi (java)

Issue:
Throw exception when writing end of an invalid header segment
Improve error message or throw an exception when a header (`ISA`, `UNB`, etc) is ended in the writer and it is not valid, as determined by the selected `Dialect`.

Pull request description:
Require a valid header segment when writing delimiter
Fixes #275

Signed-off-by: Michael Edgar <michael@xlate.io>

Patch:
diff --git a/src/main/java/io/xlate/edi/internal/stream/StaEDIStreamWriter.java b/src/main/java/io/xlate/edi/internal/stream/StaEDIStreamWriter.java
index 46af46a..5f91800 100644
--- a/src/main/java/io/xlate/edi/internal/stream/StaEDIStreamWriter.java
+++ b/src/main/java/io/xlate/edi/internal/stream/StaEDIStreamWriter.java
@@ -43,6 +43,7 @@ import io.xlate.edi.internal.stream.tokenization.EDIFACTDialect;
 import io.xlate.edi.internal.stream.tokenization.ElementDataHandler;
 import io.xlate.edi.internal.stream.tokenization.State;
 import io.xlate.edi.internal.stream.tokenization.ValidationEventHandler;
+import io.xlate.edi.internal.stream.tokenization.X12Dialect;
 import io.xlate.edi.internal.stream.validation.UsageError;
 import io.xlate.edi.internal.stream.validation.Validator;
 import io.xlate.edi.schema.EDIReference;
@@ -307,6 +308,10 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
     }
 
     private void write(int output) throws EDIStreamException {
+        write(output, false);
+    }
+
+    private void write(int output, boolean isPrettyPrint) throws EDIStreamException {
         CharacterClass clazz;
 
         clazz = characters.getClass(output);
@@ -322,7 +327,7 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
         case HEADER_EDIFACT_U: // U(NA) or U(NB)
         case HEADER_TRADACOMS_S: // S(TX)
             unconfirmedBuffer.clear();
-            writeHeader((char) output);
+            writeHeader((char) output, isPrettyPrint);
             break;
         case HEADER_X12_S:
         case HEADER_EDIFACT_N:
@@ -332,7 +337,7 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
         case HEADER_ELEMENT_END:
         case HEADER_COMPONENT_END:
         case HEADER_SEGMENT_END:
-            writeHeader((char) output);
+            writeHeader((char) output, isPrettyPrint);
             break;
         case INVALID:
             throw new EDIException(String.format("Invalid state: %s; output 0x%04X", state, output));
@@ -342,9 +347,9 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
         }
     }
 
-    void writeHeader(char output) throws EDIStreamException {
-        if (!dialect.appendHeader(characters, output)) {
-            throw new EDIStreamException(String.format("Unexpected header character: 0x%04X [%s]", (int) output, output), location);
+    void writeHeader(char output, boolean isPrettyPrint) throws EDIStreamException {
+        if (!isPrettyPrint && !dialect.appendHeader(characters, output)) {
+            throw new EDIStreamException(String.format("Failed writing %s header: %s", dialect.getStandard(), dialect.getRejectionMessage()));
         }
 
         unconfirmedBuffer.append(output);
@@ -398,6 +403,9 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
         ensureLevel(LEVEL_INITIAL);
         ensureState(State.INITIAL);
         level = LEVEL_INTERCHANGE;
+        if (controlSchema == null) {
+            LOGGER.warning("Starting interchange without control structure validation. See EDIStreamWriter#setControlSchema");
+        }
         return this;
     }
 
@@ -499,7 +507,9 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
         write(this.segmentTerminator);
 
         if (prettyPrint) {
-            writeString(prettyPrintString);
+            for (int i = 0, m = prettyPrintString.length(); i < m; i++) {
+                write(prettyPrintString.charAt(i), true);
+            }
         }
     }
 
@@ -522,6 +532,18 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
 
         writeSegmentTerminator();
 
+        switch (state) {
+        case SEGMENT_END:
+        case HEADER_SEGMENT_END:
+        case INITIAL: // Ending final segment of the interchange
+            break;
+        default:
+            if (state.isHeaderState() && dialect instanceof X12Dialect) {
+                throw new EDIStreamException("Invalid X12 ISA segment: too short or elements missing");
+            }
+            break;
+        }
+
         level = LEVEL_INTERCHANGE;
         location.clearSegmentLocations();
         transactionSchemaAllowed = false;
@@ -931,31 +953,35 @@ public class StaEDIStreamWriter implements EDIStreamWriter, ElementDataHandler,
     }
 
     private CharSequence validateElement(Runnable setupCommand, CharSequence data) {
-        return validator().map(validator -> {
-            CharSequence elementData;
+        return validator()
+                .map(validator -> validateElement(setupCommand, data, validator))
+                .orElse(data);
+    }
 
-            if (this.formatElements) {
-                elementData = this.formattedElement;
-                this.formattedElement.setLength(0);
-                this.formattedElement.append(data); // Validator will clear and re-format if configured
-            } else {
-                elementData = data;
-            }
+    CharSequence validateElement(Runnable setupCommand, CharSequence data, Validator validator) {
+        CharSequence elementData;
 
-            errors.clear();
-            setupCommand.run();
+        if (this.formatElements) {
+            elementData = this.formattedElement;
+            this.formattedElement.setLength(0);
+            this.formattedElement.append(data); // Validator will clear and re-format if configured
+        } else {
+            elementData = data;
+        }
 
-            if (!validator.validateElement(dialect, location, data, this.formattedElement)) {
-                reportElementErrors(validator, elementData);
-            }
+        errors.clear();
+        setupCommand.run();
 
-            if (!errors.isEmpty()) {
-                throw validationExceptionChain(errors);
-            }
+        if (!validator.validateElement(dialect, location, data, this.formattedElement)) {
+            reportElementErrors(validator, elementData);
+        }
+
+        if (!errors.isEmpty()) {
+            throw validationExceptionChain(errors);
+        }
 
-            dialect.elementData(elementData, location);
-            return elementData;
-        }).orElse(data);
+        dialect.elementData(elementData, location);
+        return elementData;
     }
 
     void reportElementErrors(Validator validator, CharSequence data) {
diff --git a/src/main/java/io/xlate/edi/internal/stream/tokenization/EDIFACTDialect.java b/src/main/java/io/xlate/edi/internal/stream/tokenization/EDIFACTDialect.java
index cb04d0e..d0764ac 100644
--- a/src/main/java/io/xlate/edi/internal/stream/tokenization/EDIFACTDialect.java
+++ b/src/main/java/io/xlate/edi/internal/stream/tokenization/EDIFACTDialect.java
@@ -228,7 +228,7 @@ public class EDIFACTDialect extends Dialect {
             break;
         }
 
-        if (index > EDIFACT_UNA_LENGTH) {
+        if (index >= EDIFACT_UNA_LENGTH) {
             if (characters.isIgnored(value)) {
                 header.deleteCharAt(index--);
             } else if (isIndexBeyondUNBFirstElement()) {
@@ -243,10 +243,12 @@ public class EDIFACTDialect extends Dialect {
                     unbStart = index - 2;
                 } else {
                     // Some other segment / element?
+                    rejectionMessage = String.format("Expected UNB segment following UNA but received %s%s", un, value);
                     proceed = false;
                 }
             } else if (isUnexpectedSegmentDetected(value)) {
                 // Some other segment / element?
+                rejectionMessage = String.format("Expected UNB segment following UNA but received %s", header.subSequence(EDIFACT_UNA_LENGTH, index));
                 proceed = false;
             }
         }
@@ -259,7 +261,7 @@ public class EDIFACTDialect extends Dialect {
     }
 
     boolean isUnexpectedSegmentDetected(int value) {
-        return unbStart < 0 && value == elementDelimiter;
+        return unbStart < 0 && (value == elementDelimiter || value == segmentDelimiter);
     }
 
     void setCharacterClass(CharacterSet characters, CharacterClass charClass, char value, boolean allowSpace) {
diff --git a/src/main/java/io/xlate/edi/internal/stream/tokenization/State.java b/src/main/java/io/xlate/edi/internal/stream/tokenization/State.java
index 62ff9bc..448ce6c 100644
--- a/src/main/java/io/xlate/edi/internal/stream/tokenization/State.java
+++ b/src/main/java/io/xlate/edi/internal/stream/tokenization/State.java
@@ -357,4 +357,8 @@ public enum State {
         return TRANSITIONS[dialect][code][clazz.code];
     }
 
+    public boolean isHeaderState() {
+        return Category.HEADER == code;
+    }
+
 }
diff --git a/src/main/java/io/xlate/edi/internal/stream/tokenization/X12Dialect.java b/src/main/java/io/xlate/edi/internal/stream/tokenization/X12Dialect.java
index 3fd46bd..aafc145 100644
--- a/src/main/java/io/xlate/edi/internal/stream/tokenization/X12Dialect.java
+++ b/src/main/java/io/xlate/edi/internal/stream/tokenization/X12Dialect.java
@@ -15,6 +15,10 @@
  ******************************************************************************/
 package io.xlate.edi.internal.stream.tokenization;
 
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.Set;
+
 import io.xlate.edi.stream.EDIStreamConstants.Standards;
 import io.xlate.edi.stream.Location;
 
@@ -39,7 +43,8 @@ public class X12Dialect extends Dialect {
     private static final int X12_SEGMENT_OFFSET = 105;
     private static final int X12_REPEAT_OFFSET = 82;
 
-    private static final int[] X12_ISA_TOKENS = { 3, 6, 17, 20, 31, 34, 50, 53, 69, 76, 81, 83, 89, 99, 101, 103 };
+    private static final Integer[] X12_ISA_TOKENS = { 3, 6, 17, 20, 31, 34, 50, 53, 69, 76, 81, 83, 89, 99, 101, 103 };
+    private static final Set<Integer> elementDelimiterOffsets = new HashSet<>(Arrays.asList(X12_ISA_TOKENS));
 
     private String[] version;
     char[] header;
@@ -75,13 +80,17 @@ public class X12Dialect extends Dialect {
     }
 
     boolean initialize(CharacterSet characters) {
-        final char ELEMENT = header[X12_ELEMENT_OFFSET];
-        int e = 0;
-
         for (int i = 0, m = X12_ISA_LENGTH; i < m; i++) {
-            if (ELEMENT == header[i] && X12_ISA_TOKENS[e++] != i) {
-                rejectionMessage = String.format("Unexpected element delimiter value '%s' in X12 header position %d", ELEMENT, i + 1);
-                return false;
+            if (elementDelimiterOffsets.contains(i)) {
+                if (elementDelimiter != header[i]) {
+                    rejectionMessage = String.format("Element delimiter '%s' required in position %d of X12 header but not found", elementDelimiter, i + 1);
+                    return false;
+                }
+            } else {
+                if (elementDelimiter == header[i]) {
+                    rejectionMessage = String.format("Unexpected element delimiter value '%s' in X12 header position %d", elementDelimiter, i + 1);
+                    return false;
+                }
             }
         }
 
diff --git a/src/main/java/io/xlate/edi/stream/EDIStreamWriter.java b/src/main/java/io/xlate/edi/stream/EDIStreamWriter.java
index 9b6d34e..a9953cb 100644
--- a/src/main/java/io/xlate/edi/stream/EDIStreamWriter.java
+++ b/src/main/java/io/xlate/edi/stream/EDIStreamWriter.java
@@ -69,13 +69,18 @@ public interface EDIStreamWriter extends AutoCloseable {
      * group, and transaction/message envelopes.
      * <p>
      * Calls to this method are only valid before the interchange is started.
+     * <p>
+     * A built-in control schema may be obtained from
+     * {@link io.xlate.edi.schema.SchemaFactory#getControlSchema(String, String[])
+     * SchemaFactory#getControlSchema} to pass to this method.
      *
      * @param controlSchema
      *            the schema instance to use for validation of control
      *            structures
      * @throws IllegalStateException
      *             when the writer is not in its initial state
-     *
+     * @see io.xlate.edi.schema.SchemaFactory#getControlSchema(String, String[])
+     *      SchemaFactory#getControlSchema
      * @since 1.1
      */
     void setControlSchema(Schema controlSchema);
diff --git a/src/main/resources/X12/common.xml b/src/main/resources/X12/common.xml
index ea2c500..21707d7 100644
--- a/src/main/resources/X12/common.xml
+++ b/src/main/resources/X12/common.xml
@@ -270,7 +270,7 @@
     </enumeration>
   </elementType>
 
-  <elementType name="I65" base="string" title="Repetition Separator"/>
+  <elementType name="I65" base="string" minLength="1" maxLength="1" title="Repetition Separator"/>
 
   <!-- Elements for ISX Segment (Interchange Syntax Extension) -->
   <elementType name="I69" base="string" minLength="1" maxLength="1" title="Release Character"/>


Test patch:
diff --git a/src/test/java/io/xlate/edi/internal/stream/StaEDIStreamWriterTest.java b/src/test/java/io/xlate/edi/internal/stream/StaEDIStreamWriterTest.java
index 6654192..0378db8 100644
--- a/src/test/java/io/xlate/edi/internal/stream/StaEDIStreamWriterTest.java
+++ b/src/test/java/io/xlate/edi/internal/stream/StaEDIStreamWriterTest.java
@@ -197,7 +197,7 @@ class StaEDIStreamWriterTest {
         writer.writeElement("0");
         writer.writeElement("P");
         EDIStreamException thrown = assertThrows(EDIStreamException.class, () -> writer.writeElement(":"));
-        assertEquals("Unexpected header character: 0x002A [*] in segment ISA at position 1, element 16", thrown.getMessage());
+        assertEquals("Failed writing X12 header: Element delimiter '*' required in position 18 of X12 header but not found", thrown.getMessage());
     }
 
     @Test
@@ -219,7 +219,7 @@ class StaEDIStreamWriterTest {
         writer.startInterchange();
         writer.writeStartSegment("ISA");
         writer.writeStartElement().writeElementData("E1").endElement();
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA*E1~", writer);
     }
@@ -349,8 +349,8 @@ class StaEDIStreamWriterTest {
         writer.writeStartElement()
               .startComponent()
               .endComponent()
-              .startComponent()
-              .writeEndSegment();
+              .startComponent();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA*:~", writer);
     }
@@ -369,8 +369,8 @@ class StaEDIStreamWriterTest {
               .startComponent()
               .endComponent()
               .endElement()
-              .writeEmptyElement()
-              .writeEndSegment();
+              .writeEmptyElement();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
 
         unconfirmedBufferEquals("ISA~", writer);
@@ -418,7 +418,7 @@ class StaEDIStreamWriterTest {
         writer.writeEmptyElement();
         writer.writeEmptyElement();
         writer.writeEmptyElement();
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA****~", writer);
     }
@@ -435,7 +435,7 @@ class StaEDIStreamWriterTest {
         writer.writeEmptyComponent();
         writer.writeEmptyComponent();
         writer.writeEmptyComponent();
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA*:::~", writer);
     }
@@ -454,7 +454,7 @@ class StaEDIStreamWriterTest {
         writer.writeEmptyComponent();
         writer.writeEmptyComponent();
         writer.endElement();
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA~", writer);
     }
@@ -513,7 +513,7 @@ class StaEDIStreamWriterTest {
         writer.endElement();
         writer.writeEmptyElement();
 
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
 
         unconfirmedBufferEquals("ISA**:LAST*:::LAST*:SECOND::LAST***LAST~", writer);
@@ -528,7 +528,7 @@ class StaEDIStreamWriterTest {
         writer.writeStartSegment("ISA");
         writer.writeStartElement();
         writer.writeElementData("TEST-ELEMENT");
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA*TEST-ELEMENT~", writer);
     }
@@ -554,7 +554,7 @@ class StaEDIStreamWriterTest {
         writer.writeStartSegment("ISA");
         writer.writeStartElement();
         writer.writeElementData(new char[] { 'C', 'H', 'A', 'R', 'S' }, 0, 5);
-        writer.writeEndSegment();
+        assertThrows(EDIStreamException.class, () -> writer.writeEndSegment());
         writer.flush();
         unconfirmedBufferEquals("ISA*CHARS~", writer);
     }
@@ -2150,4 +2150,25 @@ class StaEDIStreamWriterTest {
             thrown = thrown.getNextException();
         }
     }
+
+    @ParameterizedTest
+    @ValueSource(strings = { "UNX", "XYB" })
+    void testUnexpectedHeaderEDIFACT(String segmentTag) throws IOException, EDIStreamException, EDISchemaException {
+        final EDIOutputFactory factory = EDIOutputFactory.newFactory();
+        ByteArrayOutputStream stream = new ByteArrayOutputStream(4096);
+        EDIStreamException thrown = null;
+
+        try (EDIStreamWriter writer = factory.createEDIStreamWriter(stream)) {
+            writer.startInterchange();
+            writer.writeStartSegment("UNA")
+                .writeEndSegment();
+            thrown = assertThrows(EDIStreamException.class, () -> { // NOSONAR
+                // Exception thrown at different positions depending on the segment tag
+                writer.writeStartSegment(segmentTag);
+                writer.writeEndSegment();
+            });
+        }
+
+        assertEquals("Failed writing EDIFACT header: Expected UNB segment following UNA but received " + segmentTag, thrown.getMessage());
+    }
 }
diff --git a/src/test/java/io/xlate/edi/internal/stream/tokenization/X12DialectTest.java b/src/test/java/io/xlate/edi/internal/stream/tokenization/X12DialectTest.java
index f34ed33..28c82b7 100644
--- a/src/test/java/io/xlate/edi/internal/stream/tokenization/X12DialectTest.java
+++ b/src/test/java/io/xlate/edi/internal/stream/tokenization/X12DialectTest.java
@@ -17,8 +17,8 @@ package io.xlate.edi.internal.stream.tokenization;
 
 import static org.junit.jupiter.api.Assertions.assertArrayEquals;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
-import static org.junit.jupiter.api.Assertions.fail;
 
 import org.junit.jupiter.api.Test;
 
@@ -37,14 +37,19 @@ class X12DialectTest {
     }
 
     @Test
-    void testInitialize() {
-        try {
-            X12Dialect x12 = (X12Dialect) DialectFactory.getDialect("ISA".toCharArray(), 0, 3);
-            x12.header = "ISA*00*          *00*          *ZZ*ReceiverID     *ZZ*Sender         *050812*1953*^*00501*508121953*0*P*:~".toCharArray();
-            x12.initialize(new CharacterSet());
-        } catch (Exception e) {
-            fail("Unexpected exception: " + e.getMessage());
-        }
+    void testInitializeTrue() throws EDIException {
+        X12Dialect x12 = (X12Dialect) DialectFactory.getDialect("ISA".toCharArray(), 0, 3);
+        x12.header = "ISA*00*          *00*          *ZZ*ReceiverID     *ZZ*Sender         *050812*1953*^*00501*508121953*0*P*:~".toCharArray();
+        assertTrue(x12.initialize(new CharacterSet()));
+    }
+
+    @Test
+    void testInitializeRejected() throws EDIException {
+        X12Dialect x12 = (X12Dialect) DialectFactory.getDialect("ISA".toCharArray(), 0, 3);
+        x12.header = "ISA*00*    *     *00*          *ZZ*ReceiverID     *ZZ*Sender         *050812*1953*^*00501*508121953*0*P*:~".toCharArray();
+        assertFalse(x12.initialize(new CharacterSet()));
+        assertTrue(x12.isRejected());
+        assertEquals("Unexpected element delimiter value '*' in X12 header position 12", x12.getRejectionMessage());
     }
 
     @Test