import { useEffect , useRef , useState , useCallback } from 'react' ; import { useImmer } from 'use-immer' ; import { attachListeners } from './attach-listener' ; import { createSocket } from './create-socket' ; import websocketWrapper from './proxy' ; import { Options , ReadyStateState , SendMessage , WebSocketMessage , IJSON , } from './types' ; import { isValidData , getType , ReadyState } from './globals' ; import { DEFAULT_OPTIONS } from './config' ; export const useWebSocket = ( url , options = DEFAULT_OPTIONS ) : { send : SendMessage ; message : WebSocketEventMap [ 'message' ] | null ; readyState : ReadyState ; connect : ( ) => void ; wsInstance : WebSocket | null ; } => { const [ lastMessage , setLastMessage ] = useState < WebSocketEventMap [ 'message' ] | null > ( null ) ; const [ readyState , setReadyState ] = useImmer < ReadyStateState > ( { } ) ; const convertedUrl = useRef < string | null > ( null ) ; const webSocketRef = useRef < WebSocket | null > ( null ) ; const webSocketProxy = useRef < WebSocket | null > ( null ) ; const startRef = useRef < ( ) => void > ( ( ) => { } ) ; const reconnectCount = useRef < number > ( 0 ) ; const messageQueue = useRef < ( WebSocketMessage | IJSON ) [ ] > ( [ ] ) ; const expectClose = useRef < boolean > ( false ) ; const readyStateFromUrl = convertedUrl . current && readyState [ convertedUrl . current ] !== undefined ? readyState [ convertedUrl . current ] : url !== null ? ReadyState . CONNECTING : ReadyState . UNINSTANTIATED ; const sendMessage = useCallback ( message => { const type = getType ( message ) ; if ( ! isValidData ( type ) ) { console . warn ( "The data format you pass to useWebSocket's send method is invalid" ) ; return ; } const wsInstance = webSocketRef . current ; if ( wsInstance && wsInstance . readyState === ReadyState . OPEN ) { if ( type === '[object Object]' ) { wsInstance . send ( JSON . stringify ( message ) ) ; } else { wsInstance . send ( message as WebSocketMessage ) ; } } else { messageQueue . current . push ( message ) ; } } , [ ] ) ; const getWebSocket = useCallback ( ( ) => { if ( webSocketProxy . current === null && webSocketRef . current !== null ) { webSocketProxy . current = websocketWrapper ( webSocketRef . current ) ; } return webSocketProxy . current ; } , [ ] ) ; useEffect ( ( ) => { if ( url !== null ) { let removeListeners ; const start = async ( ) => { expectClose . current = false ; if ( typeof url === 'function' ) { convertedUrl . current = await url ( ) ; } else { convertedUrl . current = url ; } createSocket ( webSocketRef , convertedUrl . current , setReadyState ) ; if ( webSocketRef . current !== null ) { removeListeners = attachListeners ( webSocketRef . current , convertedUrl . current , { setLastMessage , setReadyState } , options , startRef . current , reconnectCount , expectClose , sendMessage ) ; } } ; startRef . current = ( ) => { expectClose . current = true ; if ( webSocketProxy . current ) webSocketProxy . current = null ; removeListeners ?. ( ) ; start ( ) ; } ; start ( ) ; return ( ) => { expectClose . current = true ; if ( webSocketProxy . current ) webSocketProxy . current = null ; removeListeners ?. ( ) ; } ; } } , [ url ] ) ; useEffect ( ( ) => { if ( readyStateFromUrl === ReadyState . OPEN ) { messageQueue . current . splice ( 0 ) . forEach ( message => { sendMessage ( message ) ; } ) ; } } , [ readyStateFromUrl ] ) ; return { send : sendMessage , message : lastMessage , readyState : readyStateFromUrl , connect : startRef . current , wsInstance : getWebSocket ( ) } ; } ;