/*
    * Logback: the reliable, generic, fast and flexible logging framework.
    * Copyright (C) 1999-2026, QOS.ch. All rights reserved.
    *
    * This program and the accompanying materials are dual-licensed under
    * either the terms of the Eclipse Public License v2.0 as published by
    * the Eclipse Foundation
    *
    *   or (per the licensee's choosing)
   *
   * under the terms of the GNU Lesser General Public License version 2.1
   * as published by the Free Software Foundation.
   */
  
  package ch.qos.logback.core.net;
  
  import java.io.ByteArrayInputStream;
  import java.io.ByteArrayOutputStream;
  import java.io.IOException;
  import java.io.InvalidClassException;
  import java.io.ObjectOutputStream;
  import java.util.HashSet;
  import java.util.Set;
  
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.Test;
  
  import static org.junit.jupiter.api.Assertions.assertEquals;
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertThrows;
  
  public class HardenedObjectInputStreamTest {
  
      ByteArrayOutputStream bos;
      ObjectOutputStream oos;
      HardenedObjectInputStream inputStream;
      String[] whitelist = new String[] { Innocent.class.getName() };
  
      @BeforeEach
      public void setUp() throws Exception {
          bos = new ByteArrayOutputStream();
          oos = new ObjectOutputStream(bos);
      }
  
      @AfterEach
      public void tearDown() throws Exception {
      }
  
      @Test
      public void smoke() throws ClassNotFoundException, IOException {
          Innocent innocent = new Innocent();
          innocent.setAnInt(1);
          innocent.setAnInteger(2);
          innocent.setaString("smoke");
          Innocent back = writeAndRead(innocent);
          assertEquals(innocent, back);
      }
  
      private Innocent writeAndRead(Innocent innocent) throws IOException, ClassNotFoundException {
          writeObject(oos, innocent);
          ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
          inputStream = new HardenedObjectInputStream(bis, whitelist);
          Innocent fooBack = (Innocent) inputStream.readObject();
          inputStream.close();
          return fooBack;
      }
  
      private void writeObject(ObjectOutputStream oos, Object o) throws IOException {
          oos.writeObject(o);
          oos.flush();
          oos.close();
      }
  
      @Test
      public void denialOfService() throws ClassNotFoundException, IOException {
          ByteArrayInputStream bis = new ByteArrayInputStream(payload());
          inputStream = new HardenedObjectInputStream(bis, whitelist);
          try {
              assertThrows(InvalidClassException.class, () -> inputStream.readObject());
          } finally {
              inputStream.close();
          }
      }
  
      private byte[] payload() throws IOException {
          Set root = buildEvilHashset();
          writeObject(oos, root);
          return bos.toByteArray();
      }
  
      private Set buildEvilHashset() {
          Set root = new HashSet();
          Set s1 = root;
          Set s2 = new HashSet();
          for (int i = 0; i < 100; i++) {
              Set t1 = new HashSet();
              Set t2 = new HashSet();
              t1.add("foo"); // make it not equal to t2
             s1.add(t1);
             s1.add(t2);
             s2.add(t1);
             s2.add(t2);
             s1 = t1;
             s2 = t2;
         }
         return root;
     }
 }