1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20 package org.apache.wss4j.dom;
21
22
23
24
25
26
27
28
29
30
31
32
33 import java.util.ArrayList;
34 import java.util.Collections;
35 import java.util.HashMap;
36 import java.util.LinkedList;
37 import java.util.List;
38 import java.util.Map;
39
40 import javax.xml.crypto.dom.DOMCryptoContext;
41
42 import org.apache.wss4j.common.crypto.Crypto;
43 import org.apache.wss4j.common.ext.WSSecurityException;
44 import org.apache.wss4j.common.util.XMLUtils;
45 import org.apache.wss4j.dom.callback.CallbackLookup;
46 import org.apache.wss4j.dom.engine.WSSecurityEngineResult;
47 import org.w3c.dom.Document;
48 import org.w3c.dom.Element;
49
50 public class WSDocInfo {
51 private Document doc;
52 private Crypto crypto;
53
54
55
56
57
58 private final Map<String, TokenValue> tokens = new HashMap<>();
59
60 private final List<WSSecurityEngineResult> results = new LinkedList<>();
61 private final Map<Integer, List<WSSecurityEngineResult>> actionResults = new HashMap<>();
62 private CallbackLookup callbackLookup;
63 private Element securityHeader;
64
65 public WSDocInfo(Document doc) {
66
67
68
69
70
71 if (doc != null && doc.getDocumentElement() != null) {
72 this.doc = doc.getDocumentElement().getOwnerDocument();
73 } else {
74 this.doc = doc;
75 }
76 }
77
78
79
80
81 public void clear() {
82 crypto = null;
83 doc = null;
84 callbackLookup = null;
85 securityHeader = null;
86 tokens.clear();
87 results.clear();
88 actionResults.clear();
89 }
90
91
92
93
94
95
96 public void addTokenElement(Element element) throws WSSecurityException {
97 addTokenElement(element, true);
98 }
99
100
101
102
103
104
105
106 public void addTokenElement(Element element, boolean checkMultipleElements) throws WSSecurityException {
107 if (element == null) {
108 return;
109 }
110
111 if (element.hasAttributeNS(WSConstants.WSU_NS, "Id")) {
112 String id = element.getAttributeNS(WSConstants.WSU_NS, "Id");
113 TokenValue tokenValue = new TokenValue("Id", WSConstants.WSU_NS, element);
114 TokenValue previousValue = tokens.put(id, tokenValue);
115 if (checkMultipleElements && previousValue != null) {
116 throw new WSSecurityException(
117 WSSecurityException.ErrorCode.INVALID_SECURITY_TOKEN, "duplicateError"
118 );
119 }
120 }
121
122 if (element.hasAttributeNS(null, "Id")) {
123 String id = element.getAttributeNS(null, "Id");
124 TokenValue tokenValue = new TokenValue("Id", null, element);
125 TokenValue previousValue = tokens.put(id, tokenValue);
126 if (checkMultipleElements && previousValue != null) {
127 throw new WSSecurityException(
128 WSSecurityException.ErrorCode.INVALID_SECURITY_TOKEN, "duplicateError"
129 );
130 }
131 }
132
133
134 if ("Assertion".equals(element.getLocalName())) {
135 if (WSConstants.SAML_NS.equals(element.getNamespaceURI())
136 && element.hasAttributeNS(null, "AssertionID")) {
137 String id = element.getAttributeNS(null, "AssertionID");
138 TokenValue tokenValue = new TokenValue("AssertionID", null, element);
139 TokenValue previousValue = tokens.put(id, tokenValue);
140 if (checkMultipleElements && previousValue != null) {
141 throw new WSSecurityException(
142 WSSecurityException.ErrorCode.INVALID_SECURITY_TOKEN, "duplicateError"
143 );
144 }
145 } else if (WSConstants.SAML2_NS.equals(element.getNamespaceURI())
146 && element.hasAttributeNS(null, "ID")) {
147 String id = element.getAttributeNS(null, "ID");
148 TokenValue tokenValue = new TokenValue("ID", null, element);
149 TokenValue previousValue = tokens.put(id, tokenValue);
150 if (checkMultipleElements && previousValue != null) {
151 throw new WSSecurityException(
152 WSSecurityException.ErrorCode.INVALID_SECURITY_TOKEN, "duplicateError"
153 );
154 }
155 }
156 }
157
158 }
159
160
161
162
163
164
165
166
167 public Element getTokenElement(String uri) {
168 String id = XMLUtils.getIDFromReference(uri);
169 if (id == null) {
170 return null;
171 }
172
173 TokenValue token = tokens.get(id);
174 if (token != null) {
175 return token.getToken();
176 }
177
178 return null;
179 }
180
181
182
183
184
185 public void setTokensOnContext(DOMCryptoContext context) {
186 if (!tokens.isEmpty() && context != null) {
187 for (Map.Entry<String, TokenValue> entry : tokens.entrySet()) {
188 TokenValue tokenValue = entry.getValue();
189 context.setIdAttributeNS(tokenValue.getToken(), tokenValue.getIdNamespace(),
190 tokenValue.getIdName());
191 }
192 }
193 }
194
195 public void setTokenOnContext(String uri, DOMCryptoContext context) {
196 String id = XMLUtils.getIDFromReference(uri);
197 if (id == null || context == null) {
198 return;
199 }
200
201 TokenValue tokenValue = tokens.get(id);
202 if (tokenValue != null) {
203 context.setIdAttributeNS(tokenValue.getToken(), tokenValue.getIdNamespace(),
204 tokenValue.getIdName());
205 }
206 }
207
208
209
210
211
212
213 public void addResult(WSSecurityEngineResult result) {
214 results.add(result);
215 Integer resultTag = (Integer)result.get(WSSecurityEngineResult.TAG_ACTION);
216 if (resultTag != null) {
217 List<WSSecurityEngineResult> storedResults = actionResults.get(resultTag);
218 if (storedResults == null) {
219 storedResults = new ArrayList<>();
220 }
221 storedResults.add(result);
222 actionResults.put(resultTag, storedResults);
223 }
224 }
225
226
227
228
229
230 public List<WSSecurityEngineResult> getResults() {
231 if (results.isEmpty()) {
232 return Collections.emptyList();
233 }
234 return new ArrayList<>(results);
235 }
236
237
238
239
240
241 public Map<Integer, List<WSSecurityEngineResult>> getActionResults() {
242 if (actionResults.isEmpty()) {
243 return Collections.emptyMap();
244 }
245 return new HashMap<>(actionResults);
246 }
247
248
249
250
251
252
253 public WSSecurityEngineResult getResult(String uri) {
254 String id = XMLUtils.getIDFromReference(uri);
255 if (id != null && !results.isEmpty()) {
256 for (WSSecurityEngineResult result : results) {
257 String cId = (String)result.get(WSSecurityEngineResult.TAG_ID);
258 if (id.equals(cId)) {
259 return result;
260 }
261 }
262 }
263 return null;
264 }
265
266
267
268
269 public List<WSSecurityEngineResult> getResultsByTag(Integer tag) {
270 if (actionResults.isEmpty() || !actionResults.containsKey(tag)) {
271 return Collections.emptyList();
272 }
273
274 return Collections.unmodifiableList(actionResults.get(tag));
275 }
276
277
278
279
280 public boolean hasResult(Integer tag, String uri) {
281 String id = XMLUtils.getIDFromReference(uri);
282 if (id == null || uri == null || uri.length() == 0) {
283 return false;
284 }
285
286 if (!actionResults.isEmpty() && actionResults.containsKey(tag)) {
287 for (WSSecurityEngineResult result : actionResults.get(tag)) {
288 String cId = (String)result.get(WSSecurityEngineResult.TAG_ID);
289 if (id.equals(cId)) {
290 return true;
291 }
292 }
293 }
294
295 return false;
296 }
297
298
299
300
301
302 public Crypto getCrypto() {
303 return crypto;
304 }
305
306
307
308
309 public Document getDocument() {
310 return doc;
311 }
312
313
314
315
316
317 public void setCrypto(Crypto crypto) {
318 this.crypto = crypto;
319 }
320
321
322
323
324 public void setCallbackLookup(CallbackLookup callbackLookup) {
325 this.callbackLookup = callbackLookup;
326 }
327
328
329
330
331 public CallbackLookup getCallbackLookup() {
332 return callbackLookup;
333 }
334
335
336
337
338 public Element getSecurityHeader() {
339 return securityHeader;
340 }
341
342
343
344
345
346
347 public void setSecurityHeader(Element securityHeader) {
348 this.securityHeader = securityHeader;
349 }
350
351 private static class TokenValue {
352 private final String idName;
353 private final String idNamespace;
354 private final Element token;
355
356 TokenValue(String idName, String idNamespace, Element token) {
357 this.idName = idName;
358 this.idNamespace = idNamespace;
359 this.token = token;
360 }
361
362 public String getIdName() {
363 return idName;
364 }
365
366 public String getIdNamespace() {
367 return idNamespace;
368 }
369 public Element getToken() {
370 return token;
371 }
372 }
373
374 }