View Javadoc
1   /*
2    * ====================================================================
3    * Licensed to the Apache Software Foundation (ASF) under one
4    * or more contributor license agreements.  See the NOTICE file
5    * distributed with this work for additional information
6    * regarding copyright ownership.  The ASF licenses this file
7    * to you under the Apache License, Version 2.0 (the
8    * "License"); you may not use this file except in compliance
9    * with the License.  You may obtain a copy of the License at
10   *
11   *   http://www.apache.org/licenses/LICENSE-2.0
12   *
13   * Unless required by applicable law or agreed to in writing,
14   * software distributed under the License is distributed on an
15   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16   * KIND, either express or implied.  See the License for the
17   * specific language governing permissions and limitations
18   * under the License.
19   * ====================================================================
20   *
21   * This software consists of voluntary contributions made by many
22   * individuals on behalf of the Apache Software Foundation.  For more
23   * information on the Apache Software Foundation, please see
24   * <http://www.apache.org/>.
25   *
26   */
27  package org.apache.hc.core5.http.nio.entity;
28  
29  import java.io.IOException;
30  import java.nio.ByteBuffer;
31  import java.security.MessageDigest;
32  import java.security.NoSuchAlgorithmException;
33  import java.util.ArrayList;
34  import java.util.LinkedHashSet;
35  import java.util.List;
36  import java.util.Set;
37  
38  import org.apache.hc.core5.http.Header;
39  import org.apache.hc.core5.http.message.BasicHeader;
40  import org.apache.hc.core5.http.nio.AsyncEntityProducer;
41  import org.apache.hc.core5.http.nio.DataStreamChannel;
42  import org.apache.hc.core5.util.Args;
43  import org.apache.hc.core5.util.TextUtils;
44  
45  /**
46   * @since 5.0
47   */
48  public class DigestingEntityProducer implements AsyncEntityProducer {
49  
50      private final AsyncEntityProducer wrapped;
51      private final MessageDigest digester;
52  
53      private volatile byte[] digest;
54  
55      public DigestingEntityProducer(
56              final String algo,
57              final AsyncEntityProducer wrapped) {
58          this.wrapped = Args.notNull(wrapped, "Entity consumer");
59          try {
60              this.digester = MessageDigest.getInstance(algo);
61          } catch (final NoSuchAlgorithmException ex) {
62              throw new IllegalArgumentException("Unsupported digest algorithm: " + algo);
63          }
64      }
65  
66      @Override
67      public boolean isRepeatable() {
68          return wrapped.isRepeatable();
69      }
70  
71      @Override
72      public long getContentLength() {
73          return wrapped.getContentLength();
74      }
75  
76      @Override
77      public String getContentType() {
78          return wrapped.getContentType();
79      }
80  
81      @Override
82      public String getContentEncoding() {
83          return wrapped.getContentEncoding();
84      }
85  
86      @Override
87      public boolean isChunked() {
88          return wrapped.isChunked();
89      }
90  
91      @Override
92      public Set<String> getTrailerNames() {
93          final Set<String> allNames = new LinkedHashSet<>();
94          final Set<String> names = wrapped.getTrailerNames();
95          if (names != null) {
96              allNames.addAll(names);
97          }
98          allNames.add("digest-algo");
99          allNames.add("digest");
100         return allNames;
101     }
102 
103     @Override
104     public int available() {
105         return wrapped.available();
106     }
107 
108     @Override
109     public void produce(final DataStreamChannel channel) throws IOException {
110         wrapped.produce(new DataStreamChannel() {
111 
112             @Override
113             public void requestOutput() {
114                 channel.requestOutput();
115             }
116 
117             @Override
118             public int write(final ByteBuffer src) throws IOException {
119                 final ByteBuffer dup = src.duplicate();
120                 final int writtenBytes = channel.write(src);
121                 if (writtenBytes > 0) {
122                     dup.limit(dup.position() + writtenBytes);
123                     digester.update(dup);
124                 }
125                 return writtenBytes;
126             }
127 
128             @Override
129             public void endStream(final List<? extends Header> trailers) throws IOException {
130                 digest = digester.digest();
131                 final List<Header> allTrailers = new ArrayList<>();
132                 if (trailers != null) {
133                     allTrailers.addAll(trailers);
134                 }
135                 allTrailers.add(new BasicHeader("digest-algo", digester.getAlgorithm()));
136                 allTrailers.add(new BasicHeader("digest", TextUtils.toHexString(digest)));
137                 channel.endStream(allTrailers);
138             }
139 
140             @Override
141             public void endStream() throws IOException {
142                 endStream(null);
143             }
144 
145         });
146     }
147 
148     @Override
149     public void failed(final Exception cause) {
150         wrapped.failed(cause);
151     }
152 
153     @Override
154     public void releaseResources() {
155         wrapped.releaseResources();
156     }
157 
158     public byte[] getDigest() {
159         return digest;
160     }
161 
162 }