1 /*
2  * JBoss, Home of Professional Open Source.
3  * Copyright 2014 Red Hat, Inc., and individual contributors
4  * as indicated by the @author tags.
5  *
6  * Licensed under the Apache License, Version 2.0 (the "License");
7  * you may not use this file except in compliance with the License.
8  * You may obtain a copy of the License at
9  *
10  *     http://www.apache.org/licenses/LICENSE-2.0
11  *
12  *  Unless required by applicable law or agreed to in writing, software
13  *  distributed under the License is distributed on an "AS IS" BASIS,
14  *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15  *  See the License for the specific language governing permissions and
16  *  limitations under the License.
17  */
18
19 package io.undertow.util;
20
21 import java.util.AbstractCollection;
22 import java.util.Arrays;
23 import java.util.Collection;
24 import java.util.Collections;
25 import java.util.Iterator;
26 import java.util.NoSuchElementException;
27
28 /**
29  * An optimized array-backed header map.
30  *
31  * @author <a href="mailto:david.lloyd@redhat.com">David M. Lloyd</a>
32  */
33 public final class HeaderMap implements Iterable<HeaderValues> {
34
35     private Object[] table;
36     private int size;
37     private Collection<HttpString> headerNames;
38
39     public HeaderMap() {
40         table = new Object[16];
41     }
42
43     private HeaderValues getEntry(final HttpString headerName) {
44         if (headerName == null) {
45             return null;
46         }
47         final int hc = headerName.hashCode();
48         final int idx = hc & (table.length - 1);
49         final Object o = table[idx];
50         if (o == null) {
51             return null;
52         }
53         HeaderValues headerValues;
54         if (o instanceof HeaderValues) {
55             headerValues = (HeaderValues) o;
56             if (! headerName.equals(headerValues.key)) {
57                 return null;
58             }
59             return headerValues;
60         } else {
61             final HeaderValues[] row = (HeaderValues[]) o;
62             for (int i = 0; i < row.length; i++) {
63                 headerValues = row[i];
64                 if (headerValues != null && headerName.equals(headerValues.key)) {
65                     return headerValues;
66                 }
67             }
68             return null;
69         }
70     }
71
72
73     private HeaderValues getEntry(final String headerName) {
74         if (headerName == null) {
75             return null;
76         }
77         final int hc = HttpString.hashCodeOf(headerName);
78         final int idx = hc & (table.length - 1);
79         final Object o = table[idx];
80         if (o == null) {
81             return null;
82         }
83         HeaderValues headerValues;
84         if (o instanceof HeaderValues) {
85             headerValues = (HeaderValues) o;
86             if (! headerValues.key.equalToString(headerName)) {
87                 return null;
88             }
89             return headerValues;
90         } else {
91             final HeaderValues[] row = (HeaderValues[]) o;
92             for (int i = 0; i < row.length; i++) {
93                 headerValues = row[i];
94                 if (headerValues != null && headerValues.key.equalToString(headerName)) {
95                     return headerValues;
96                 }
97             }
98             return null;
99         }
100     }
101
102     private HeaderValues removeEntry(final HttpString headerName) {
103         if (headerName == null) {
104             return null;
105         }
106         final int hc = headerName.hashCode();
107         final Object[] table = this.table;
108         final int idx = hc & (table.length - 1);
109         final Object o = table[idx];
110         if (o == null) {
111             return null;
112         }
113         HeaderValues headerValues;
114         if (o instanceof HeaderValues) {
115             headerValues = (HeaderValues) o;
116             if (! headerName.equals(headerValues.key)) {
117                 return null;
118             }
119             table[idx] = null;
120             size --;
121             return headerValues;
122         } else {
123             final HeaderValues[] row = (HeaderValues[]) o;
124             for (int i = 0; i < row.length; i++) {
125                 headerValues = row[i];
126                 if (headerValues != null && headerName.equals(headerValues.key)) {
127                     row[i] = null;
128                     size --;
129                     return headerValues;
130                 }
131             }
132             return null;
133         }
134     }
135
136
137     private HeaderValues removeEntry(final String headerName) {
138         if (headerName == null) {
139             return null;
140         }
141         final int hc = HttpString.hashCodeOf(headerName);
142         final Object[] table = this.table;
143         final int idx = hc & (table.length - 1);
144         final Object o = table[idx];
145         if (o == null) {
146             return null;
147         }
148         HeaderValues headerValues;
149         if (o instanceof HeaderValues) {
150             headerValues = (HeaderValues) o;
151             if (! headerValues.key.equalToString(headerName)) {
152                 return null;
153             }
154             table[idx] = null;
155             size --;
156             return headerValues;
157         } else {
158             final HeaderValues[] row = (HeaderValues[]) o;
159             for (int i = 0; i < row.length; i++) {
160                 headerValues = row[i];
161                 if (headerValues != null && headerValues.key.equalToString(headerName)) {
162                     row[i] = null;
163                     size --;
164                     return headerValues;
165                 }
166             }
167             return null;
168         }
169     }
170
171     private void resize() {
172         final int oldLen = table.length;
173         if (oldLen == 0x40000000) {
174             return;
175         }
176         assert Integer.bitCount(oldLen) == 1;
177         Object[] newTable = Arrays.copyOf(table, oldLen << 1);
178         table = newTable;
179         for (int i = 0; i < oldLen; i ++) {
180             if (newTable[i] == null) {
181                 continue;
182             }
183             if (newTable[i] instanceof HeaderValues) {
184                 HeaderValues e = (HeaderValues) newTable[i];
185                 if ((e.key.hashCode() & oldLen) != 0) {
186                     newTable[i] = null;
187                     newTable[i + oldLen] = e;
188                 }
189                 continue;
190             }
191             HeaderValues[] oldRow = (HeaderValues[]) newTable[i];
192             HeaderValues[] newRow = oldRow.clone();
193             int rowLen = oldRow.length;
194             newTable[i + oldLen] = newRow;
195             HeaderValues item;
196             for (int j = 0; j < rowLen; j ++) {
197                 item = oldRow[j];
198                 if (item != null) {
199                     if ((item.key.hashCode() & oldLen) != 0) {
200                         oldRow[j] = null;
201                     } else {
202                         newRow[j] = null;
203                     }
204                 }
205             }
206         }
207     }
208
209     private HeaderValues getOrCreateEntry(final HttpString headerName) {
210         if (headerName == null) {
211             return null;
212         }
213         final int hc = headerName.hashCode();
214         final Object[] table = this.table;
215         final int length = table.length;
216         final int idx = hc & (length - 1);
217         final Object o = table[idx];
218         HeaderValues headerValues;
219         if (o == null) {
220             if (size >= length >> 1) {
221                 resize();
222                 return getOrCreateEntry(headerName);
223             }
224             headerValues = new HeaderValues(headerName);
225             table[idx] = headerValues;
226             size++;
227             return headerValues;
228         }
229         return getOrCreateNonEmpty(headerName, table, length, idx, o);
230     }
231
232     private HeaderValues getOrCreateNonEmpty(HttpString headerName, Object[] table, int length, int idx, Object o) {
233         HeaderValues headerValues;
234         if (o instanceof HeaderValues) {
235             headerValues = (HeaderValues) o;
236             if (! headerName.equals(headerValues.key)) {
237                 if (size >= length >> 1) {
238                     resize();
239                     return getOrCreateEntry(headerName);
240                 }
241                 size++;
242                 final HeaderValues[] row = { headerValues, new HeaderValues(headerName), null, null };
243                 table[idx] = row;
244                 return row[1];
245             }
246             return headerValues;
247         } else {
248             final HeaderValues[] row = (HeaderValues[]) o;
249             int empty = -1;
250             for (int i = 0; i < row.length; i++) {
251                 headerValues = row[i];
252                 if (headerValues != null) {
253                     if (headerName.equals(headerValues.key)) {
254                         return headerValues;
255                     }
256                 } else if (empty == -1) {
257                     empty = i;
258                 }
259             }
260             if (size >= length >> 1) {
261                 resize();
262                 return getOrCreateEntry(headerName);
263             }
264             size++;
265             headerValues = new HeaderValues(headerName);
266             if (empty != -1) {
267                 row[empty] = headerValues;
268             } else {
269                 if (row.length >= 16) {
270                     throw new SecurityException("Excessive collisions");
271                 }
272                 final HeaderValues[] newRow = Arrays.copyOf(row, row.length + 3);
273                 newRow[row.length] = headerValues;
274                 table[idx] = newRow;
275             }
276             return headerValues;
277         }
278     }
279
280     // get
281
282     public HeaderValues get(final HttpString headerName) {
283         return getEntry(headerName);
284     }
285
286     public HeaderValues get(final String headerName) {
287         return getEntry(headerName);
288     }
289
290     public String getFirst(HttpString headerName) {
291         HeaderValues headerValues = getEntry(headerName);
292         if (headerValues == null) return null;
293         return headerValues.getFirst();
294     }
295
296     public String getFirst(String headerName) {
297         HeaderValues headerValues = getEntry(headerName);
298         if (headerValues == null) return null;
299         return headerValues.getFirst();
300     }
301
302     public String get(HttpString headerName, int index) throws IndexOutOfBoundsException {
303         if (headerName == null) {
304             return null;
305         }
306         final HeaderValues headerValues = getEntry(headerName);
307         if (headerValues == null) {
308             return null;
309         }
310         return headerValues.get(index);
311     }
312
313     public String get(String headerName, int index) throws IndexOutOfBoundsException {
314         if (headerName == null) {
315             return null;
316         }
317         final HeaderValues headerValues = getEntry(headerName);
318         if (headerValues == null) {
319             return null;
320         }
321         return headerValues.get(index);
322     }
323
324     public String getLast(HttpString headerName) {
325         if (headerName == null) {
326             return null;
327         }
328         HeaderValues headerValues = getEntry(headerName);
329         if (headerValues == null) return null;
330         return headerValues.getLast();
331     }
332
333     public String getLast(String headerName) {
334         if (headerName == null) {
335             return null;
336         }
337         HeaderValues headerValues = getEntry(headerName);
338         if (headerValues == null) return null;
339         return headerValues.getLast();
340     }
341
342     // count
343
344     public int count(HttpString headerName) {
345         if (headerName == null) {
346             return 0;
347         }
348         final HeaderValues headerValues = getEntry(headerName);
349         if (headerValues == null) {
350             return 0;
351         }
352         return headerValues.size();
353     }
354
355     public int count(String headerName) {
356         if (headerName == null) {
357             return 0;
358         }
359         final HeaderValues headerValues = getEntry(headerName);
360         if (headerValues == null) {
361             return 0;
362         }
363         return headerValues.size();
364     }
365
366     public int size() {
367         return size;
368     }
369
370     // iterate
371
372     /**
373      * Do a fast iteration of this header map without creating any objects.
374      *
375      * @return an opaque iterating cookie, or -1 if no iteration is possible
376      *
377      * @see #fiNext(long)
378      * @see #fiCurrent(long)
379      */
380     public long fastIterate() {
381         final Object[] table = this.table;
382         final int len = table.length;
383         int ri = 0;
384         int ci;
385         while (ri < len) {
386             final Object item = table[ri];
387             if (item != null) {
388                 if (item instanceof HeaderValues) {
389                     return (long)ri << 32L;
390                 } else {
391                     final HeaderValues[] row = (HeaderValues[]) item;
392                     ci = 0;
393                     final int rowLen = row.length;
394                     while (ci < rowLen) {
395                         if (row[ci] != null) {
396                             return (long)ri << 32L | (ci & 0xffffffffL);
397                         }
398                         ci ++;
399                     }
400                 }
401             }
402             ri++;
403         }
404         return -1L;
405     }
406
407     /**
408      * Do a fast iteration of this header map without creating any objects, only considering non-empty header values.
409      *
410      * @return an opaque iterating cookie, or -1 if no iteration is possible
411      */
412     public long fastIterateNonEmpty() {
413         final Object[] table = this.table;
414         final int len = table.length;
415         int ri = 0;
416         int ci;
417         while (ri < len) {
418             final Object item = table[ri];
419             if (item != null) {
420                 if (item instanceof HeaderValues) {
421                     if(!((HeaderValues) item).isEmpty()) {
422                         return (long) ri << 32L;
423                     }
424                 } else {
425                     final HeaderValues[] row = (HeaderValues[]) item;
426                     ci = 0;
427                     final int rowLen = row.length;
428                     while (ci < rowLen) {
429                         if (row[ci] != null && !row[ci].isEmpty()) {
430                             return (long)ri << 32L | (ci & 0xffffffffL);
431                         }
432                         ci ++;
433                     }
434                 }
435             }
436             ri++;
437         }
438         return -1L;
439     }
440
441     /**
442      * Find the next index in a fast iteration.
443      *
444      * @param cookie the previous cookie value
445      * @return the next cookie value, or -1L if iteration is done
446      */
447     public long fiNext(long cookie) {
448         if (cookie == -1L) return -1L;
449         final Object[] table = this.table;
450         final int len = table.length;
451         int ri = (int) (cookie >> 32);
452         int ci = (int) cookie;
453         Object item = table[ri];
454         if (item instanceof HeaderValues[]) {
455             final HeaderValues[] row = (HeaderValues[]) item;
456             final int rowLen = row.length;
457             if (++ci >= rowLen) {
458                 ri ++; ci = 0;
459             } else if (row[ci] != null) {
460                 return (long)ri << 32L | (ci & 0xffffffffL);
461             }
462         } else {
463             ri ++; ci = 0;
464         }
465         while (ri < len) {
466             item = table[ri];
467             if (item instanceof HeaderValues) {
468                 return (long)ri << 32L;
469             } else if (item instanceof HeaderValues[]) {
470                 final HeaderValues[] row = (HeaderValues[]) item;
471                 final int rowLen = row.length;
472                 while (ci < rowLen) {
473                     if (row[ci] != null) {
474                         return (long)ri << 32L | (ci & 0xffffffffL);
475                     }
476                     ci ++;
477                 }
478             }
479             ci = 0;
480             ri ++;
481         }
482         return -1L;
483     }
484
485     /**
486      * Find the next non-empty index in a fast iteration.
487      *
488      * @param cookie the previous cookie value
489      * @return the next cookie value, or -1L if iteration is done
490      */
491     public long fiNextNonEmpty(long cookie) {
492         if (cookie == -1L) return -1L;
493         final Object[] table = this.table;
494         final int len = table.length;
495         int ri = (int) (cookie >> 32);
496         int ci = (int) cookie;
497         Object item = table[ri];
498         if (item instanceof HeaderValues[]) {
499             final HeaderValues[] row = (HeaderValues[]) item;
500             final int rowLen = row.length;
501             if (++ci >= rowLen) {
502                 ri ++; ci = 0;
503             } else if (row[ci] != null && !row[ci].isEmpty()) {
504                 return (long)ri << 32L | (ci & 0xffffffffL);
505             }
506         } else {
507             ri ++; ci = 0;
508         }
509         while (ri < len) {
510             item = table[ri];
511             if (item instanceof HeaderValues && !((HeaderValues) item).isEmpty()) {
512                 return (long)ri << 32L;
513             } else if (item instanceof HeaderValues[]) {
514                 final HeaderValues[] row = (HeaderValues[]) item;
515                 final int rowLen = row.length;
516                 while (ci < rowLen) {
517                     if (row[ci] != null && !row[ci].isEmpty()) {
518                         return (long)ri << 32L | (ci & 0xffffffffL);
519                     }
520                     ci ++;
521                 }
522             }
523             ci = 0;
524             ri ++;
525         }
526         return -1L;
527     }
528
529     /**
530      * Return the value at the current index in a fast iteration.
531      *
532      * @param cookie the iteration cookie value
533      * @return the values object at this position
534      * @throws NoSuchElementException if the cookie value is invalid
535      */
536     public HeaderValues fiCurrent(long cookie) {
537         try {
538             final Object[] table = this.table;
539             int ri = (int) (cookie >> 32);
540             int ci = (int) cookie;
541             final Object item = table[ri];
542             if (item instanceof HeaderValues[]) {
543                 return ((HeaderValues[])item)[ci];
544             } else if (ci == 0) {
545                 return (HeaderValues) item;
546             } else {
547                 throw new NoSuchElementException();
548             }
549         } catch (RuntimeException e) {
550             throw new NoSuchElementException();
551         }
552     }
553
554     public Iterable<String> eachValue(final HttpString headerName) {
555         if (headerName == null) {
556             return Collections.emptyList();
557         }
558         final HeaderValues entry = getEntry(headerName);
559         if (entry == null) {
560             return Collections.emptyList();
561         }
562         return entry;
563     }
564
565     public Iterator<HeaderValues> iterator() {
566         return new Iterator<HeaderValues>() {
567             final Object[] table = HeaderMap.this.table;
568             boolean consumed;
569             int ri, ci;
570
571             private HeaderValues _next() {
572                 for (;;) {
573                     if (ri >= table.length) {
574                         return null;
575                     }
576                     final Object o = table[ri];
577                     if (o == null) {
578                         // zero-entry row
579                         ri++;
580                         ci = 0;
581                         consumed = false;
582                         continue;
583                     }
584                     if (o instanceof HeaderValues) {
585                         // one-entry row
586                         if (ci > 0 || consumed) {
587                             ri++;
588                             ci = 0;
589                             consumed = false;
590                             continue;
591                         }
592                         return (HeaderValues) o;
593                     }
594                     final HeaderValues[] row = (HeaderValues[]) o;
595                     final int len = row.length;
596                     if (ci >= len) {
597                         ri ++;
598                         ci = 0;
599                         consumed = false;
600                         continue;
601                     }
602                     if (consumed) {
603                         ci++;
604                         consumed = false;
605                         continue;
606                     }
607                     final HeaderValues headerValues = row[ci];
608                     if (headerValues == null) {
609                         ci ++;
610                         continue;
611                     }
612                     return headerValues;
613                 }
614             }
615
616             public boolean hasNext() {
617                 return _next() != null;
618             }
619
620             public HeaderValues next() {
621                 final HeaderValues next = _next();
622                 if (next == null) {
623                     throw new NoSuchElementException();
624                 }
625                 consumed = true;
626                 return next;
627             }
628
629             public void remove() {
630             }
631         };
632     }
633
634     public Collection<HttpString> getHeaderNames() {
635         if (headerNames != null) {
636             return headerNames;
637         }
638         return headerNames = new AbstractCollection<HttpString>() {
639             public boolean contains(final Object o) {
640                 return o instanceof HttpString && getEntry((HttpString) o) != null;
641             }
642
643             public boolean add(final HttpString httpString) {
644                 getOrCreateEntry(httpString);
645                 return true;
646             }
647
648             public boolean remove(final Object o) {
649                 if (! (o instanceof HttpString)) return false;
650                 HttpString s = (HttpString) o;
651                 HeaderValues entry = getEntry(s);
652                 if (entry == null) {
653                     return false;
654                 }
655                 entry.clear();
656                 return true;
657             }
658
659             public void clear() {
660                 HeaderMap.this.clear();
661             }
662
663             public Iterator<HttpString> iterator() {
664                 final Iterator<HeaderValues> iterator = HeaderMap.this.iterator();
665                 return new Iterator<HttpString>() {
666                     public boolean hasNext() {
667                         return iterator.hasNext();
668                     }
669
670                     public HttpString next() {
671                         return iterator.next().getHeaderName();
672                     }
673
674                     public void remove() {
675                         iterator.remove();
676                     }
677                 };
678             }
679
680             public int size() {
681                 return HeaderMap.this.size();
682             }
683         };
684     }
685
686     // add
687
688     public HeaderMap add(HttpString headerName, String headerValue) {
689         addLast(headerName, headerValue);
690         return this;
691     }
692
693     public HeaderMap addFirst(final HttpString headerName, final String headerValue) {
694         if (headerName == null) {
695             throw new IllegalArgumentException("headerName is null");
696         }
697         if (headerValue == null) {
698             return this;
699         }
700         getOrCreateEntry(headerName).addFirst(headerValue);
701         return this;
702     }
703
704     public HeaderMap addLast(final HttpString headerName, final String headerValue) {
705         if (headerName == null) {
706             throw new IllegalArgumentException("headerName is null");
707         }
708         if (headerValue == null) {
709             return this;
710         }
711         getOrCreateEntry(headerName).addLast(headerValue);
712         return this;
713     }
714
715     public HeaderMap add(HttpString headerName, long headerValue) {
716         add(headerName, Long.toString(headerValue));
717         return this;
718     }
719
720
721     public HeaderMap addAll(HttpString headerName, Collection<String> headerValues) {
722         if (headerName == null) {
723             throw new IllegalArgumentException("headerName is null");
724         }
725         if (headerValues == null || headerValues.isEmpty()) {
726             return this;
727         }
728         getOrCreateEntry(headerName).addAll(headerValues);
729         return this;
730     }
731
732     // put
733
734     public HeaderMap put(HttpString headerName, String headerValue) {
735         if (headerName == null) {
736             throw new IllegalArgumentException("headerName is null");
737         }
738         if (headerValue == null) {
739             remove(headerName);
740             return this;
741         }
742         final HeaderValues headerValues = getOrCreateEntry(headerName);
743         headerValues.clear();
744         headerValues.add(headerValue);
745         return this;
746     }
747
748     public HeaderMap put(HttpString headerName, long headerValue) {
749         if (headerName == null) {
750             throw new IllegalArgumentException("headerName is null");
751         }
752         final HeaderValues entry = getOrCreateEntry(headerName);
753         entry.clear();
754         entry.add(Long.toString(headerValue));
755         return this;
756     }
757
758     public HeaderMap putAll(HttpString headerName, Collection<String> headerValues) {
759         if (headerName == null) {
760             throw new IllegalArgumentException("headerName is null");
761         }
762         if (headerValues == null || headerValues.isEmpty()) {
763             remove(headerName);
764             return this;
765         }
766         final HeaderValues entry = getOrCreateEntry(headerName);
767         entry.clear();
768         entry.addAll(headerValues);
769         return this;
770     }
771
772     // clear
773
774     public HeaderMap clear() {
775         Arrays.fill(table, null);
776         size = 0;
777         return this;
778     }
779
780     // remove
781
782     public Collection<String> remove(HttpString headerName) {
783         if (headerName == null) {
784             return Collections.emptyList();
785         }
786         final Collection<String> values = removeEntry(headerName);
787         return values != null ? values : Collections.<String>emptyList();
788     }
789
790     public Collection<String> remove(String headerName) {
791         if (headerName == null) {
792             return Collections.emptyList();
793         }
794         final Collection<String> values = removeEntry(headerName);
795         return values != null ? values : Collections.<String>emptyList();
796     }
797
798     // contains
799
800     public boolean contains(HttpString headerName) {
801         final HeaderValues headerValues = getEntry(headerName);
802         if (headerValues == null) {
803             return false;
804         }
805         final Object v = headerValues.value;
806         if (v instanceof String) {
807             return true;
808         }
809         final String[] list = (String[]) v;
810         for (int i = 0; i < list.length; i++) {
811             if (list[i] != null) {
812                 return true;
813             }
814         }
815         return false;
816     }
817
818     public boolean contains(String headerName) {
819         final HeaderValues headerValues = getEntry(headerName);
820         if (headerValues == null) {
821             return false;
822         }
823         final Object v = headerValues.value;
824         if (v instanceof String) {
825             return true;
826         }
827         final String[] list = (String[]) v;
828         for (int i = 0; i < list.length; i++) {
829             if (list[i] != null) {
830                 return true;
831             }
832         }
833         return false;
834     }
835
836     // compare
837
838     @Override
839     public boolean equals(final Object o) {
840         return o == this;
841     }
842
843     @Override
844     public int hashCode() {
845         return super.hashCode();
846     }
847
848     @Override
849     public String toString() {
850         StringBuilder sb = new StringBuilder("{");
851         boolean first = true;
852         for(HttpString name : getHeaderNames()) {
853             if(first) {
854                 first = false;
855             } else {
856                 sb.append(", ");
857             }
858             sb.append(name);
859             sb.append("=[");
860             boolean f = true;
861             for(String val : get(name)) {
862                 if(f) {
863                     f = false;
864                 } else {
865                     sb.append(", ");
866                 }
867                 sb.append(val);
868             }
869             sb.append("]");
870         }
871         sb.append("}");
872         return sb.toString();
873     }
874 }
875