package com.xmzs.midjourney.wss.user;
|
|
import cn.hutool.core.exceptions.ValidateException;
|
import cn.hutool.core.text.CharSequenceUtil;
|
import cn.hutool.core.thread.ThreadUtil;
|
import cn.hutool.core.util.RandomUtil;
|
import com.xmzs.midjourney.ProxyProperties;
|
import com.xmzs.midjourney.ReturnCode;
|
import com.xmzs.midjourney.domain.DiscordAccount;
|
import com.xmzs.midjourney.util.AsyncLockUtils;
|
import com.xmzs.midjourney.wss.WebSocketStarter;
|
import com.neovisionaries.ws.client.WebSocket;
|
import com.neovisionaries.ws.client.WebSocketAdapter;
|
import com.neovisionaries.ws.client.WebSocketFactory;
|
import com.neovisionaries.ws.client.WebSocketFrame;
|
import eu.bitwalker.useragentutils.UserAgent;
|
import lombok.extern.slf4j.Slf4j;
|
import net.dv8tion.jda.api.utils.data.DataArray;
|
import net.dv8tion.jda.api.utils.data.DataObject;
|
import net.dv8tion.jda.api.utils.data.DataType;
|
import net.dv8tion.jda.internal.requests.WebSocketCode;
|
import net.dv8tion.jda.internal.utils.compress.Decompressor;
|
import net.dv8tion.jda.internal.utils.compress.ZlibDecompressor;
|
|
import java.nio.charset.StandardCharsets;
|
import java.time.Duration;
|
import java.util.List;
|
import java.util.Map;
|
import java.util.concurrent.Executors;
|
import java.util.concurrent.Future;
|
import java.util.concurrent.ScheduledExecutorService;
|
import java.util.concurrent.TimeUnit;
|
import java.util.concurrent.TimeoutException;
|
|
@Slf4j
|
public class UserWebSocketStarter extends WebSocketAdapter implements WebSocketStarter {
|
private static final int CONNECT_RETRY_LIMIT = 3;
|
|
private final ProxyProperties.ProxyConfig proxyConfig;
|
private final DiscordAccount account;
|
private final UserMessageListener userMessageListener;
|
private final ScheduledExecutorService heartExecutor;
|
private final String wssServer;
|
private final DataObject authData;
|
|
private Decompressor decompressor;
|
private WebSocket socket = null;
|
private String resumeGatewayUrl;
|
private String sessionId;
|
|
private Future<?> heartbeatInterval;
|
private Future<?> heartbeatTimeout;
|
private boolean heartbeatAck = false;
|
private Object sequence = null;
|
private long interval = 41250;
|
private boolean trying = false;
|
|
public UserWebSocketStarter(String wssServer, DiscordAccount account, UserMessageListener userMessageListener, ProxyProperties.ProxyConfig proxyConfig) {
|
this.wssServer = wssServer;
|
this.account = account;
|
this.userMessageListener = userMessageListener;
|
this.proxyConfig = proxyConfig;
|
this.heartExecutor = Executors.newSingleThreadScheduledExecutor();
|
this.authData = createAuthData();
|
}
|
|
@Override
|
public void setTrying(boolean trying) {
|
this.trying = trying;
|
}
|
|
@Override
|
public synchronized void start() throws Exception {
|
this.decompressor = new ZlibDecompressor(2048);
|
WebSocketFactory webSocketFactory = createWebSocketFactory(this.proxyConfig);
|
String gatewayUrl = CharSequenceUtil.isNotBlank(this.resumeGatewayUrl) ? this.resumeGatewayUrl : this.wssServer;
|
this.socket = webSocketFactory.createSocket(gatewayUrl + "/?encoding=json&v=9&compress=zlib-stream");
|
this.socket.addListener(this);
|
this.socket.addHeader("Accept-Encoding", "gzip, deflate, br")
|
.addHeader("Accept-Language", "zh-CN,zh;q=0.9")
|
.addHeader("Cache-Control", "no-cache")
|
.addHeader("Pragma", "no-cache")
|
.addHeader("Sec-Websocket-Extensions", "permessage-deflate; client_max_window_bits")
|
.addHeader("User-Agent", this.account.getUserAgent());
|
this.socket.connect();
|
}
|
|
@Override
|
public void onConnected(WebSocket websocket, Map<String, List<String>> headers) {
|
log.debug("[wss-{}] Connected to websocket.", this.account.getDisplay());
|
}
|
|
@Override
|
public void handleCallbackError(WebSocket websocket, Throwable cause) throws Exception {
|
log.error("[wss-{}] There was some websocket error.", this.account.getDisplay(), cause);
|
}
|
|
@Override
|
public void onDisconnected(WebSocket websocket, WebSocketFrame serverCloseFrame, WebSocketFrame clientCloseFrame, boolean closedByServer) throws Exception {
|
int code;
|
String closeReason;
|
if (closedByServer) {
|
code = serverCloseFrame.getCloseCode();
|
closeReason = serverCloseFrame.getCloseReason();
|
} else {
|
code = clientCloseFrame.getCloseCode();
|
closeReason = clientCloseFrame.getCloseReason();
|
}
|
connectFinish(code, closeReason);
|
if (this.trying) {
|
return;
|
}
|
if (code == 5240) {
|
// 隐式关闭wss
|
clearAllStates();
|
} else if (code >= 4000) {
|
log.warn("[wss-{}] Can't reconnect! Account disabled. Closed by {}({}).", this.account.getDisplay(), code, closeReason);
|
clearAllStates();
|
this.account.setEnable(false);
|
} else if (code == 2001) {
|
// reconnect
|
log.warn("[wss-{}] Waiting try reconnect...", this.account.getDisplay());
|
tryReconnect();
|
} else {
|
log.warn("[wss-{}] Closed by {}({}). Waiting try new connection...", this.account.getDisplay(), code, closeReason);
|
tryNewConnect();
|
}
|
}
|
|
private void tryReconnect() {
|
clearSocketStates();
|
try {
|
this.trying = true;
|
tryStart(true);
|
} catch (Exception e) {
|
if (e instanceof TimeoutException) {
|
sendClose(5240, "try new connect");
|
}
|
log.warn("[wss-{}] Try reconnect fail: {}, Waiting try new connection...", this.account.getDisplay(), e.getMessage());
|
ThreadUtil.sleep(1000);
|
tryNewConnect();
|
}
|
}
|
|
private void tryNewConnect() {
|
this.trying = true;
|
for (int i = 1; i <= CONNECT_RETRY_LIMIT; i++) {
|
clearAllStates();
|
try {
|
tryStart(false);
|
return;
|
} catch (Exception e) {
|
if (e instanceof TimeoutException) {
|
sendClose(5240, "try new connect");
|
}
|
log.warn("[wss-{}] Try new connection fail ({}): {}", this.account.getDisplay(), i, e.getMessage());
|
ThreadUtil.sleep(5000);
|
}
|
}
|
log.error("[wss-{}] Account disabled", this.account.getDisplay());
|
this.account.setEnable(false);
|
}
|
|
public void tryStart(boolean reconnect) throws Exception {
|
start();
|
AsyncLockUtils.LockObject lock = AsyncLockUtils.waitForLock("wss:" + this.account.getChannelId(), Duration.ofSeconds(20));
|
int code = lock.getProperty("code", Integer.class, 0);
|
if (code == ReturnCode.SUCCESS) {
|
log.debug("[wss-{}] {} success.", this.account.getDisplay(), reconnect ? "Reconnect" : "New connect");
|
return;
|
}
|
throw new ValidateException(lock.getProperty("description", String.class));
|
}
|
|
@Override
|
public void onBinaryMessage(WebSocket websocket, byte[] binary) throws Exception {
|
if (this.decompressor == null) {
|
return;
|
}
|
byte[] decompressBinary = this.decompressor.decompress(binary);
|
if (decompressBinary == null) {
|
return;
|
}
|
String json = new String(decompressBinary, StandardCharsets.UTF_8);
|
DataObject data = DataObject.fromJson(json);
|
int opCode = data.getInt("op");
|
switch (opCode) {
|
case WebSocketCode.HEARTBEAT -> {
|
log.debug("[wss-{}] Receive heartbeat.", this.account.getDisplay());
|
handleHeartbeat();
|
}
|
case WebSocketCode.HEARTBEAT_ACK -> {
|
this.heartbeatAck = true;
|
clearHeartbeatTimeout();
|
}
|
case WebSocketCode.HELLO -> {
|
handleHello(data);
|
doResumeOrIdentify();
|
}
|
case WebSocketCode.RESUME -> {
|
log.debug("[wss-{}] Receive resumed.", this.account.getDisplay());
|
connectSuccess();
|
}
|
case WebSocketCode.RECONNECT -> sendReconnect("receive server reconnect");
|
case WebSocketCode.INVALIDATE_SESSION -> sendClose(1009, "receive session invalid");
|
case WebSocketCode.DISPATCH -> handleDispatch(data);
|
default -> log.debug("[wss-{}] Receive unknown code: {}.", this.account.getDisplay(), data);
|
}
|
}
|
|
private void handleHello(DataObject data) {
|
clearHeartbeatInterval();
|
this.interval = data.getObject("d").getLong("heartbeat_interval");
|
this.heartbeatAck = true;
|
this.heartbeatInterval = this.heartExecutor.scheduleAtFixedRate(() -> {
|
if (this.heartbeatAck) {
|
this.heartbeatAck = false;
|
send(WebSocketCode.HEARTBEAT, this.sequence);
|
} else {
|
sendReconnect("heartbeat has not ack interval");
|
}
|
}, (long) Math.floor(RandomUtil.randomDouble(0, 1) * this.interval), this.interval, TimeUnit.MILLISECONDS);
|
}
|
|
private void doResumeOrIdentify() {
|
if (CharSequenceUtil.isBlank(this.sessionId)) {
|
log.debug("[wss-{}] Send identify msg.", this.account.getDisplay());
|
send(WebSocketCode.IDENTIFY, this.authData);
|
} else {
|
log.debug("[wss-{}] Send resume msg.", this.account.getDisplay());
|
send(WebSocketCode.RESUME, DataObject.empty().put("token", this.account.getUserToken())
|
.put("session_id", this.sessionId).put("seq", this.sequence));
|
}
|
}
|
|
private void handleHeartbeat() {
|
send(WebSocketCode.HEARTBEAT, this.sequence);
|
this.heartbeatTimeout = ThreadUtil.execAsync(() -> {
|
ThreadUtil.sleep(this.interval);
|
sendReconnect("heartbeat has not ack");
|
});
|
}
|
|
private void clearAllStates() {
|
clearSocketStates();
|
clearResumeStates();
|
}
|
|
private void clearSocketStates() {
|
clearHeartbeatTimeout();
|
clearHeartbeatInterval();
|
this.socket = null;
|
this.decompressor = null;
|
}
|
|
private void clearResumeStates() {
|
this.sessionId = null;
|
this.sequence = null;
|
this.resumeGatewayUrl = null;
|
}
|
|
private void clearHeartbeatInterval() {
|
if (this.heartbeatInterval != null) {
|
this.heartbeatInterval.cancel(true);
|
this.heartbeatInterval = null;
|
}
|
}
|
|
private void clearHeartbeatTimeout() {
|
if (this.heartbeatTimeout != null) {
|
this.heartbeatTimeout.cancel(true);
|
this.heartbeatTimeout = null;
|
}
|
}
|
|
private void sendReconnect(String reason) {
|
sendClose(2001, reason);
|
}
|
|
private void sendClose(int code, String reason) {
|
if (this.socket != null) {
|
this.socket.sendClose(code, reason);
|
}
|
}
|
|
private void send(int op, Object d) {
|
if (this.socket != null) {
|
this.socket.sendText(DataObject.empty().put("op", op).put("d", d).toString());
|
}
|
}
|
|
private void connectSuccess() {
|
this.trying = false;
|
connectFinish(ReturnCode.SUCCESS, "");
|
}
|
|
private void connectFinish(int code, String description) {
|
AsyncLockUtils.LockObject lock = AsyncLockUtils.getLock("wss:" + this.account.getChannelId());
|
if (lock != null) {
|
lock.setProperty("code", code);
|
lock.setProperty("description", description);
|
lock.awake();
|
}
|
}
|
|
private void handleDispatch(DataObject raw) {
|
this.sequence = raw.opt("s").orElse(null);
|
if (!raw.isType("d", DataType.OBJECT)) {
|
return;
|
}
|
DataObject content = raw.getObject("d");
|
String t = raw.getString("t", null);
|
if ("READY".equals(t)) {
|
this.sessionId = content.getString("session_id");
|
this.resumeGatewayUrl = content.getString("resume_gateway_url");
|
log.debug("[wss-{}] Dispatch ready: identify.", this.account.getDisplay());
|
connectSuccess();
|
return;
|
} else if ("RESUMED".equals(t)) {
|
log.debug("[wss-{}] Dispatch read: resumed.", this.account.getDisplay());
|
connectSuccess();
|
return;
|
}
|
try {
|
this.userMessageListener.onMessage(raw);
|
} catch (Exception e) {
|
log.error("[wss-{}] Handle message error", this.account.getDisplay(), e);
|
}
|
}
|
|
private DataObject createAuthData() {
|
UserAgent agent = UserAgent.parseUserAgentString(this.account.getUserAgent());
|
DataObject connectionProperties = DataObject.empty()
|
.put("browser", agent.getBrowser().getGroup().getName())
|
.put("browser_user_agent", this.account.getUserAgent())
|
.put("browser_version", agent.getBrowserVersion().toString())
|
.put("client_build_number", 222963)
|
.put("client_event_source", null)
|
.put("device", "")
|
.put("os", agent.getOperatingSystem().getName())
|
.put("referer", "https://www.midjourney.com")
|
.put("referrer_current", "")
|
.put("referring_domain", "www.midjourney.com")
|
.put("referring_domain_current", "")
|
.put("release_channel", "stable")
|
.put("system_locale", "zh-CN");
|
DataObject presence = DataObject.empty()
|
.put("activities", DataArray.empty())
|
.put("afk", false)
|
.put("since", 0)
|
.put("status", "online");
|
DataObject clientState = DataObject.empty()
|
.put("api_code_version", 0)
|
.put("guild_versions", DataObject.empty())
|
.put("highest_last_message_id", "0")
|
.put("private_channels_version", "0")
|
.put("read_state_version", 0)
|
.put("user_guild_settings_version", -1)
|
.put("user_settings_version", -1);
|
return DataObject.empty()
|
.put("capabilities", 16381)
|
.put("client_state", clientState)
|
.put("compress", false)
|
.put("presence", presence)
|
.put("properties", connectionProperties)
|
.put("token", this.account.getUserToken());
|
}
|
|
}
|
