Skip to content

Commit 210be9c

Browse files
rstoyanchevrstoyanchev
committed
Add PrincipalMessageArgumentResolver
1 parent d3cecfc commit 210be9c

8 files changed

Lines changed: 146 additions & 52 deletions

File tree

‎spring-messaging/src/main/java/org/springframework/messaging/handler/method/InvalidMessageMethodParameterException.java‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
/*
2+
* Copyright 2002-2013 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.messaging.handler.method;
18+
19+
import org.springframework.core.MethodParameter;
20+
import org.springframework.messaging.Message;
21+
import org.springframework.messaging.MessagingException;
22+
23+
24+
/**
25+
*
26+
* @author Rossen Stoyanchev
27+
* @since 4.0
28+
*/
29+
public class InvalidMessageMethodParameterException extends MessagingException {
30+
31+
private static final long serialVersionUID = -6905878930083523161L;
32+
33+
private final MethodParameter parameter;
34+
35+
36+
public InvalidMessageMethodParameterException(Message<?> message, String description,
37+
MethodParameter parameter, Throwable cause) {
38+
super(message, description, cause);
39+
this.parameter = parameter;
40+
}
41+
42+
public InvalidMessageMethodParameterException(Message<?> message, String description,
43+
MethodParameter parameter) {
44+
45+
super(message, description);
46+
this.parameter = parameter;
47+
}
48+
49+
50+
public MethodParameter getParameter() {
51+
return this.parameter;
52+
}
53+
54+
}

‎spring-messaging/src/main/java/org/springframework/messaging/simp/SimpMessageHeaderAccessor.java‎

Lines changed: 14 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616

1717
package org.springframework.messaging.simp;
1818

19+
import java.security.Principal;
1920
import java.util.Arrays;
2021
import java.util.List;
2122
import java.util.Map;
2223

2324
import org.springframework.http.MediaType;
2425
import org.springframework.messaging.Message;
26+
import org.springframework.messaging.MessageHeaders;
2527
import org.springframework.messaging.support.NativeMessageHeaderAccessor;
2628
import org.springframework.util.Assert;
2729
import org.springframework.util.CollectionUtils;
@@ -43,9 +45,6 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
4345

4446
public static final String DESTINATIONS = "destinations";
4547

46-
// TODO
47-
public static final String CONTENT_TYPE = "contentType";
48-
4948
public static final String MESSAGE_TYPE = "messageType";
5049

5150
public static final String PROTOCOL_MESSAGE_TYPE = "protocolMessageType";
@@ -54,6 +53,8 @@ public class SimpMessageHeaderAccessor extends NativeMessageHeaderAccessor {
5453

5554
public static final String SUBSCRIPTION_ID = "subscriptionId";
5655

56+
public static final String USER = "user";
57+
5758

5859
/**
5960
* A constructor for creating new message headers.
@@ -140,12 +141,11 @@ public void setDestinations(List<String> destinations) {
140141
}
141142

142143
public MediaType getContentType() {
143-
return (MediaType) getHeader(CONTENT_TYPE);
144+
return (MediaType) getHeader(MessageHeaders.CONTENT_TYPE);
144145
}
145146

146147
public void setContentType(MediaType contentType) {
147-
Assert.notNull(contentType, "contentType is required");
148-
setHeader(CONTENT_TYPE, contentType);
148+
setHeader(MessageHeaders.CONTENT_TYPE, contentType);
149149
}
150150

151151
public String getSubscriptionId() {
@@ -164,4 +164,12 @@ public void setSessionId(String sessionId) {
164164
setHeader(SESSION_ID, sessionId);
165165
}
166166

167+
public Principal getUser() {
168+
return (Principal) getHeader(USER);
169+
}
170+
171+
public void setUser(Principal principal) {
172+
setHeader(USER, principal);
173+
}
174+
167175
}

‎spring-messaging/src/main/java/org/springframework/messaging/simp/annotation/support/PrincipalMessageArgumentResolver.java‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
/*
2+
* Copyright 2002-2013 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.messaging.simp.annotation.support;
18+
19+
import java.security.Principal;
20+
21+
import org.springframework.core.MethodParameter;
22+
import org.springframework.messaging.Message;
23+
import org.springframework.messaging.handler.method.InvalidMessageMethodParameterException;
24+
import org.springframework.messaging.handler.method.MessageArgumentResolver;
25+
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
26+
27+
28+
/**
29+
* @author Rossen Stoyanchev
30+
* @since 4.0
31+
*/
32+
public class PrincipalMessageArgumentResolver implements MessageArgumentResolver {
33+
34+
35+
@Override
36+
public boolean supportsParameter(MethodParameter parameter) {
37+
Class<?> paramType = parameter.getParameterType();
38+
return Principal.class.isAssignableFrom(paramType);
39+
}
40+
41+
@Override
42+
public Object resolveArgument(MethodParameter parameter, Message<?> message) throws Exception {
43+
SimpMessageHeaderAccessor headers = SimpMessageHeaderAccessor.wrap(message);
44+
Principal user = headers.getUser();
45+
if (user == null) {
46+
throw new InvalidMessageMethodParameterException(message, "User not available", parameter);
47+
}
48+
return user;
49+
}
50+
51+
}

‎spring-messaging/src/main/java/org/springframework/messaging/simp/handler/AnnotationSimpMessageHandler.java‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,15 +37,16 @@
3737
import org.springframework.messaging.handler.annotation.MessageMapping;
3838
import org.springframework.messaging.handler.annotation.support.MessageBodyArgumentResolver;
3939
import org.springframework.messaging.handler.annotation.support.MessageExceptionHandlerMethodResolver;
40-
import org.springframework.messaging.handler.method.MessageArgumentResolverComposite;
4140
import org.springframework.messaging.handler.method.InvocableMessageHandlerMethod;
41+
import org.springframework.messaging.handler.method.MessageArgumentResolverComposite;
4242
import org.springframework.messaging.handler.method.MessageReturnValueHandlerComposite;
43-
import org.springframework.messaging.simp.annotation.SubscribeEvent;
44-
import org.springframework.messaging.simp.annotation.UnsubscribeEvent;
45-
import org.springframework.messaging.simp.annotation.support.MessageSendingReturnValueHandler;
4643
import org.springframework.messaging.simp.MessageHolder;
4744
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
4845
import org.springframework.messaging.simp.SimpMessageType;
46+
import org.springframework.messaging.simp.annotation.SubscribeEvent;
47+
import org.springframework.messaging.simp.annotation.UnsubscribeEvent;
48+
import org.springframework.messaging.simp.annotation.support.MessageSendingReturnValueHandler;
49+
import org.springframework.messaging.simp.annotation.support.PrincipalMessageArgumentResolver;
4950
import org.springframework.messaging.support.converter.MessageConverter;
5051
import org.springframework.stereotype.Controller;
5152
import org.springframework.util.Assert;
@@ -113,6 +114,7 @@ public void afterPropertiesSet() {
113114

114115
initHandlerMethods();
115116

117+
this.argumentResolvers.addResolver(new PrincipalMessageArgumentResolver());
116118
this.argumentResolvers.addResolver(new MessageBodyArgumentResolver(this.messageConverter));
117119

118120
this.returnValueHandlers.addHandler(

‎spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompMessageConverter.java‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ public class StompMessageConverter {
4646
/**
4747
* @param stompContent a complete STOMP message (without the trailing 0x00) as byte[] or String.
4848
*/
49-
public Message<?> toMessage(Object stompContent, String sessionId) {
49+
public Message<?> toMessage(Object stompContent) {
5050

5151
byte[] byteContent = null;
5252
if (stompContent instanceof String) {
@@ -91,12 +91,10 @@ else if (stompContent instanceof byte[]){
9191
}
9292
}
9393

94-
StompHeaderAccessor stompHeaders = StompHeaderAccessor.create(command, headers);
95-
stompHeaders.setSessionId(sessionId);
96-
9794
byte[] payload = new byte[totalLength - payloadIndex];
9895
System.arraycopy(byteContent, payloadIndex, payload, 0, totalLength - payloadIndex);
9996

97+
StompHeaderAccessor stompHeaders = StompHeaderAccessor.create(command, headers);
10098
return MessageBuilder.withPayload(payload).copyHeaders(stompHeaders.toMap()).build();
10199
}
102100

‎spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompRelayMessageHandler.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,9 @@
2929
import org.springframework.context.SmartLifecycle;
3030
import org.springframework.messaging.Message;
3131
import org.springframework.messaging.MessageChannel;
32-
import org.springframework.messaging.simp.handler.AbstractSimpMessageHandler;
3332
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
3433
import org.springframework.messaging.simp.SimpMessageType;
34+
import org.springframework.messaging.simp.handler.AbstractSimpMessageHandler;
3535
import org.springframework.messaging.support.MessageBuilder;
3636
import org.springframework.util.Assert;
3737
import org.springframework.util.StringUtils;
@@ -350,7 +350,7 @@ private void readStompFrame(String stompFrame) {
350350
return;
351351
}
352352

353-
Message<?> message = stompMessageConverter.toMessage(stompFrame, this.sessionId);
353+
Message<?> message = stompMessageConverter.toMessage(stompFrame);
354354
if (logger.isTraceEnabled()) {
355355
logger.trace("Reading message " + message);
356356
}
@@ -369,6 +369,10 @@ private void readStompFrame(String stompFrame) {
369369
}
370370
relaySessions.remove(this.sessionId);
371371
}
372+
373+
headers.setSessionId(this.sessionId);
374+
message = MessageBuilder.fromMessage(message).copyHeaders(headers.toMap()).build();
375+
372376
sendMessageToClient(message);
373377
}
374378

‎spring-messaging/src/main/java/org/springframework/messaging/simp/stomp/StompWebSocketHandler.java‎

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,11 @@ public void afterConnectionEstablished(WebSocketSession session) throws Exceptio
8181
protected void handleTextMessage(WebSocketSession session, TextMessage textMessage) {
8282
try {
8383
String payload = textMessage.getPayload();
84-
Message<?> message = this.stompMessageConverter.toMessage(payload, session.getId());
84+
Message<?> message = this.stompMessageConverter.toMessage(payload);
85+
86+
StompHeaderAccessor headers = StompHeaderAccessor.wrap(message);
87+
headers.setSessionId(session.getId());
88+
headers.setUser(session.getPrincipal());
8589

8690
// TODO: validate size limits
8791
// http://stomp.github.io/stomp-specification-1.2.html#Size_Limits
@@ -96,18 +100,8 @@ protected void handleTextMessage(WebSocketSession session, TextMessage textMessa
96100
if (SimpMessageType.CONNECT.equals(messageType)) {
97101
handleConnect(session, message);
98102
}
99-
else if (SimpMessageType.MESSAGE.equals(messageType)) {
100-
handlePublish(message);
101-
}
102-
else if (SimpMessageType.SUBSCRIBE.equals(messageType)) {
103-
handleSubscribe(message);
104-
}
105-
else if (SimpMessageType.UNSUBSCRIBE.equals(messageType)) {
106-
handleUnsubscribe(message);
107-
}
108-
else if (SimpMessageType.DISCONNECT.equals(messageType)) {
109-
handleDisconnect(message);
110-
}
103+
104+
message = MessageBuilder.fromMessage(message).copyHeaders(headers.toMap()).build();
111105
this.outputChannel.send(message);
112106
}
113107
catch (Throwable t) {
@@ -124,7 +118,7 @@ else if (SimpMessageType.DISCONNECT.equals(messageType)) {
124118
}
125119
}
126120

127-
protected void handleConnect(final WebSocketSession session, Message<?> message) throws IOException {
121+
protected void handleConnect(WebSocketSession session, Message<?> message) throws IOException {
128122

129123
StompHeaderAccessor connectHeaders = StompHeaderAccessor.wrap(message);
130124
StompHeaderAccessor connectedHeaders = StompHeaderAccessor.create(StompCommand.CONNECTED);
@@ -152,18 +146,6 @@ else if (acceptVersions.isEmpty()) {
152146
session.sendMessage(new TextMessage(new String(bytes, Charset.forName("UTF-8"))));
153147
}
154148

155-
protected void handlePublish(Message<?> stompMessage) {
156-
}
157-
158-
protected void handleSubscribe(Message<?> message) {
159-
}
160-
161-
protected void handleUnsubscribe(Message<?> message) {
162-
}
163-
164-
protected void handleDisconnect(Message<?> message) {
165-
}
166-
167149
protected void sendErrorMessage(WebSocketSession session, Throwable error) {
168150

169151
StompHeaderAccessor headers = StompHeaderAccessor.create(StompCommand.ERROR);

‎spring-messaging/src/test/java/org/springframework/messaging/simp/stomp/StompMessageConverterTests.java‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -24,9 +24,6 @@
2424
import org.springframework.messaging.MessageHeaders;
2525
import org.springframework.messaging.simp.SimpMessageHeaderAccessor;
2626
import org.springframework.messaging.simp.SimpMessageType;
27-
import org.springframework.messaging.simp.stomp.StompCommand;
28-
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
29-
import org.springframework.messaging.simp.stomp.StompMessageConverter;
3027

3128
import static org.junit.Assert.*;
3229

@@ -51,17 +48,16 @@ public void connectFrame() throws Exception {
5148
String accept = "accept-version:1.1\n";
5249
String host = "host:github.org\n";
5350
String frame = "\n\n\nCONNECT\n" + accept + host + "\n";
54-
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(frame.getBytes("UTF-8"), "session-123");
51+
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(frame.getBytes("UTF-8"));
5552

5653
assertEquals(0, message.getPayload().length);
5754

5855
MessageHeaders headers = message.getHeaders();
5956
StompHeaderAccessor stompHeaders = StompHeaderAccessor.wrap(message);
6057
Map<String, Object> map = stompHeaders.toMap();
61-
assertEquals(6, map.size());
58+
assertEquals(5, map.size());
6259
assertNotNull(map.get(MessageHeaders.ID));
6360
assertNotNull(map.get(MessageHeaders.TIMESTAMP));
64-
assertNotNull(map.get(SimpMessageHeaderAccessor.SESSION_ID));
6561
assertNotNull(map.get(SimpMessageHeaderAccessor.NATIVE_HEADERS));
6662
assertNotNull(map.get(SimpMessageHeaderAccessor.MESSAGE_TYPE));
6763
assertNotNull(map.get(SimpMessageHeaderAccessor.PROTOCOL_MESSAGE_TYPE));
@@ -71,7 +67,6 @@ public void connectFrame() throws Exception {
7167

7268
assertEquals(SimpMessageType.CONNECT, stompHeaders.getMessageType());
7369
assertEquals(StompCommand.CONNECT, stompHeaders.getStompCommand());
74-
assertEquals("session-123", stompHeaders.getSessionId());
7570
assertNotNull(headers.get(MessageHeaders.ID));
7671
assertNotNull(headers.get(MessageHeaders.TIMESTAMP));
7772

@@ -89,7 +84,7 @@ public void connectWithEscapes() throws Exception {
8984
String host = "ho\\c\\ns\\rt:st\\nomp.gi\\cthu\\b.org\n";
9085
String frame = "CONNECT\n" + accept + host + "\n";
9186
@SuppressWarnings("unchecked")
92-
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(frame.getBytes("UTF-8"), "session-123");
87+
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(frame.getBytes("UTF-8"));
9388

9489
assertEquals(0, message.getPayload().length);
9590

@@ -111,7 +106,7 @@ public void connectCR12() throws Exception {
111106
String host = "host:github.org\n";
112107
String test = "CONNECT\r\n" + accept.replaceAll("\n", "\r\n") + host.replaceAll("\n", "\r\n") + "\r\n";
113108
@SuppressWarnings("unchecked")
114-
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(test.getBytes("UTF-8"), "session-123");
109+
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(test.getBytes("UTF-8"));
115110

116111
assertEquals(0, message.getPayload().length);
117112

@@ -133,7 +128,7 @@ public void connectWithEscapesAndCR12() throws Exception {
133128
String host = "ho\\c\\ns\\rt:st\\nomp.gi\\cthu\\b.org\n";
134129
String test = "\n\n\nCONNECT\r\n" + accept.replaceAll("\n", "\r\n") + host.replaceAll("\n", "\r\n") + "\r\n";
135130
@SuppressWarnings("unchecked")
136-
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(test.getBytes("UTF-8"), "session-123");
131+
Message<byte[]> message = (Message<byte[]>) this.converter.toMessage(test.getBytes("UTF-8"));
137132

138133
assertEquals(0, message.getPayload().length);
139134

0 commit comments

Comments
 (0)