From 313a7836b091d20dbb9763b3677d559fea32d13d Mon Sep 17 00:00:00 2001
From: Rossen Stoyanchev <rstoyanchev@pivotal.io>
Date: Thu, 9 Jul 2020 15:42:28 +0300
Subject: [PATCH] Provide access to Netty's WebSocketFrame

Closes gh-25099
---
 .../web/reactive/socket/WebSocketMessage.java | 32 ++++++++++++++++++-
 .../adapter/NettyWebSocketSessionSupport.java |  4 +--
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketMessage.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketMessage.java
index 832eb8d048..7d2d3d8136 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketMessage.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/WebSocketMessage.java
@@ -39,6 +39,9 @@ public class WebSocketMessage {
 
 	private final DataBuffer payload;
 
+	@Nullable
+	private final Object nativeMessage;
+
 
 	/**
 	 * Constructor for a WebSocketMessage.
@@ -47,13 +50,25 @@ public class WebSocketMessage {
 	 * then invoke this constructor.
 	 */
 	public WebSocketMessage(Type type, DataBuffer payload) {
+		this(type, payload, null);
+	}
+
+	/**
+	 * Constructor for an inbound message with access to the underlying message.
+	 * @param type the type of WebSocket message
+	 * @param payload the message content
+	 * @param nativeMessage the message from the API of the underlying WebSocket
+	 * library, if applicable.
+	 * @since 5.3
+	 */
+	public WebSocketMessage(Type type, DataBuffer payload, @Nullable Object nativeMessage) {
 		Assert.notNull(type, "'type' must not be null");
 		Assert.notNull(payload, "'payload' must not be null");
 		this.type = type;
 		this.payload = payload;
+		this.nativeMessage = nativeMessage;
 	}
 
-
 	/**
 	 * Return the message type (text, binary, etc).
 	 */
@@ -68,6 +83,21 @@ public class WebSocketMessage {
 		return this.payload;
 	}
 
+	/**
+	 * Return the message from the API of the underlying WebSocket library. This
+	 * is applicable for inbound messages only and when the underlying message
+	 * has additional fields other than the content. Currently this is the case
+	 * for Reactor Netty only.
+	 * @param <T> the type to cast the underlying message to
+	 * @return the underlying message, or {@code null}
+	 * @since 5.3
+	 */
+	@Nullable
+	@SuppressWarnings("unchecked")
+	public <T> T getNativeMessage() {
+		return (T) this.nativeMessage;
+	}
+
 	/**
 	 * A variant of {@link #getPayloadAsText(Charset)} that uses {@code UTF-8}
 	 * for decoding the raw content to text.
diff --git a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/NettyWebSocketSessionSupport.java b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/NettyWebSocketSessionSupport.java
index 5936fe7bbf..5185b401c7 100644
--- a/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/NettyWebSocketSessionSupport.java
+++ b/spring-webflux/src/main/java/org/springframework/web/reactive/socket/adapter/NettyWebSocketSessionSupport.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2018 the original author or authors.
+ * Copyright 2002-2020 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -74,7 +74,7 @@ public abstract class NettyWebSocketSessionSupport<T> extends AbstractWebSocketS
 
 	protected WebSocketMessage toMessage(WebSocketFrame frame) {
 		DataBuffer payload = bufferFactory().wrap(frame.content());
-		return new WebSocketMessage(messageTypes.get(frame.getClass()), payload);
+		return new WebSocketMessage(messageTypes.get(frame.getClass()), payload, frame);
 	}
 
 	protected WebSocketFrame toFrame(WebSocketMessage message) {