content: /*
   * Copyright The OpenTelemetry Authors
   * SPDX-License-Identifier: Apache-2.0
   */
  
  package io.opentelemetry.javaagent.instrumentation.servlet.snippet;
  
  import static java.util.logging.Level.FINE;
  
  import io.opentelemetry.javaagent.bootstrap.servlet.InjectionState;
  import java.io.IOException;
  import java.io.OutputStream;
  import java.io.UnsupportedEncodingException;
  import java.util.function.Supplier;
  import java.util.logging.Logger;
  
  public class OutputStreamSnippetInjectionHelper {
  
    private static final Logger logger =
        Logger.getLogger(OutputStreamSnippetInjectionHelper.class.getName());
  
    private final Supplier<String> snippetSupplier;
  
    public OutputStreamSnippetInjectionHelper(Supplier<String> snippetSupplier) {
      this.snippetSupplier = snippetSupplier;
    }
  
    public String getSnippet() {
      return snippetSupplier.get();
    }
  
    /**
     * return true means this method performed the injection, return false means it didn't inject
     * anything Servlet3OutputStreamWriteAdvice would skip the write method when the return value is
     * true, and would write the original bytes when the return value is false.
     */
    public boolean handleWrite(
        InjectionState state, OutputStream out, byte[] original, int off, int length)
        throws IOException {
      if (state.isHeadTagWritten()) {
        return false;
      }
      int endOfHeadTagPosition;
      boolean endOfHeadTagFound = false;
      for (endOfHeadTagPosition = off;
          endOfHeadTagPosition < length && endOfHeadTagPosition - off < length;
          endOfHeadTagPosition++) {
        if (state.processByte(original[endOfHeadTagPosition])) {
          endOfHeadTagFound = true;
          break;
        }
      }
      if (!endOfHeadTagFound) {
        return false;
      }
  
      if (state.getWrapper().isNotSafeToInject()) {
        return false;
      }
      byte[] snippetBytes;
      try {
        snippetBytes = snippetSupplier.get().getBytes(state.getCharacterEncoding());
      } catch (UnsupportedEncodingException e) {
        logger.log(FINE, "Failed getting snippet bytes", e);
        return false;
      }
      // updating Content-Length before any further writing in case that writing triggers a flush
      state.getWrapper().updateContentLengthIfPreviouslySet();
      out.write(original, off, endOfHeadTagPosition + 1);
      out.write(snippetBytes);
      out.write(original, endOfHeadTagPosition + 1, length - endOfHeadTagPosition - 1);
      return true;
    }
  
    public boolean handleWrite(InjectionState state, OutputStream out, int b) throws IOException {
      if (state.isHeadTagWritten()) {
        return false;
      }
      if (!state.processByte(b)) {
        return false;
      }
  
      if (state.getWrapper().isNotSafeToInject()) {
        return false;
      }
      byte[] snippetBytes;
      try {
        snippetBytes = snippetSupplier.get().getBytes(state.getCharacterEncoding());
      } catch (UnsupportedEncodingException e) {
        logger.log(FINE, "Failed getting snippet bytes", e);
        return false;
      }
      state.getWrapper().updateContentLengthIfPreviouslySet();
      out.write(b);
  
      out.write(snippetBytes);
      return true;
    }
  }
  