feat: expose inner WebSocketStream

[0xJepsen]

• I have looked for existing issues (including closed) about this

Feature Request

I would like to be able to get the inner WebSocketStream from the WebSocket object.

I prepose a small simple inner() function to be implemented on the WebSocket that Consumes self and get the inner [tokio_tungstenite::WebSocketStream].

Motivation

I would like to use ws_stream_tungstenite::ws_stream::new() with this object to use https://github.com/tlsnotary/tlsn/blob/3554db83e17b2e5fc98293b397a2907b7f023496/notary/server/src/service/axum_websocket.rs#L363 to create proofs of tls sessions.

Proposal

The feature could be a three line PR with the following logic in the following location. If need be it can be behind a feature flag.

  pub fn inner(self) -> WebSocketStream<TokioIo<hyper::upgrade::Upgraded>> {
    self.inner
  }

The implementation can go here:

axum/axum/src/extract/ws.rs

Lines 463 to 488 in 5027472

	impl WebSocket {
	/// Receive another message.
	///
	/// Returns `None` if the stream has closed.
	pub async fn recv(&mut self) -> Option<Result<Message, Error>> {
	self.next().await
	}
	/// Send a message.
	pub async fn send(&mut self, msg: Message) -> Result<(), Error> {
	self.inner
	.send(msg.into_tungstenite())
	.await
	.map_err(Error::new)
	}
	/// Gracefully close this WebSocket.
	pub async fn close(mut self) -> Result<(), Error> {
	self.inner.close(None).await.map_err(Error::new)
	}
	/// Return the selected WebSocket subprotocol, if one has been chosen.
	pub fn protocol(&self) -> Option<&HeaderValue> {
	self.protocol.as_ref()
	}
	}

Alternatives

I have looked into implementing this downstream in ws_stream_tungstenite but they need access to the ATungSocket<S> to check the buffer size and then call TungWebSocket::new() If another route forward is viable that i missed Im open to hearing about it and applying it to my needs.

This change would also enable the tlsn to not need to fork the ws.rs file with the changes here

https://github.com/tlsnotary/tlsn/blob/3554db83e17b2e5fc98293b397a2907b7f023496/notary/server/src/service/axum_websocket.rs#L363, which would enable closing: tlsnotary/tlsn#541

[jplatte]

The tungstenite types are intentionally not exposed by axum::extract::ws. See axum-tungstenite for a crate that supports this use case.