1use crate::{
2 config::{BanManagerConfig, ConnectionConfig, DifficultyConfig},
3 id_manager::IDManager,
4 router::Router,
5 types::ReadyIndicator,
6 BanManager, Config, ConfigManager, Result, SessionList, StratumServer,
7};
8use extended_primitives::Buffer;
9use std::{marker::PhantomData, sync::Arc};
10use tokio::{net::TcpListener, task::JoinSet};
11use tokio_stream::wrappers::TcpListenerStream;
12use tokio_util::sync::CancellationToken;
13
14#[derive(Default)]
15pub struct StratumServerBuilder<State, CState> {
16 pub server_id: u8,
17 pub host: String,
18 pub port: u16,
19 #[cfg(feature = "api")]
20 pub api_host: String,
21 #[cfg(feature = "api")]
22 pub api_port: u16,
23 pub connection_config: ConnectionConfig,
24 pub var_diff_config: DifficultyConfig,
25 pub state: State,
26 pub connection_state: PhantomData<CState>,
27 pub ready_indicator: ReadyIndicator,
28 pub shutdown_message: Option<Buffer>,
29 pub cancel_token: Option<CancellationToken>,
30 pub ban_manager_enabled: bool,
31}
32
33impl<State: Clone + Send + Sync + 'static, CState: Default + Clone + Send + Sync + 'static>
34 StratumServerBuilder<State, CState>
35{
36 pub fn new(state: State, server_id: u8) -> Self {
37 Self {
38 server_id,
39 host: String::new(),
40 port: 0,
41 #[cfg(feature = "api")]
42 api_host: String::from("0.0.0.0"),
43 #[cfg(feature = "api")]
44 api_port: 8888,
45 connection_config: ConnectionConfig::default(),
46 state,
47 connection_state: PhantomData,
48 ready_indicator: ReadyIndicator::new(false),
49 var_diff_config: DifficultyConfig {
50 retarget_share_amount: 30,
51 initial_difficulty: 16384,
52 var_diff: false,
53 minimum_difficulty: 64,
54 maximum_difficulty: 4_611_686_018_427_387_904,
55 retarget_time: 300,
56 target_time: 10,
57 variance_percent: 30.0,
58 },
59 shutdown_message: None,
65 cancel_token: None,
66 ban_manager_enabled: false,
67 }
68 }
69
70 #[must_use]
71 pub fn with_host(mut self, host: &str) -> Self {
72 self.host = host.to_owned();
73 self
74 }
75
76 #[must_use]
77 pub fn with_port(mut self, port: u16) -> Self {
78 self.port = port;
79 self
80 }
81
82 #[cfg(feature = "api")]
83 #[must_use]
84 pub fn with_api_host(mut self, host: &str) -> Self {
85 self.api_host = host.to_owned();
86 self
87 }
88
89 #[cfg(feature = "api")]
90 #[must_use]
91 pub fn with_api_port(mut self, port: u16) -> Self {
92 self.api_port = port;
93 self
94 }
95
96 #[must_use]
97 pub fn with_max_connections(mut self, max_connections: usize) -> Self {
98 self.connection_config.max_connections = Some(max_connections);
99 self
100 }
101
102 #[must_use]
103 pub fn with_proxy(mut self, value: bool) -> Self {
104 self.connection_config.proxy_protocol = value;
105 self
106 }
107
108 #[must_use]
109 pub fn with_var_diff(mut self, value: bool) -> Self {
110 self.var_diff_config.var_diff = value;
111 self
112 }
113
114 #[must_use]
115 pub fn with_minimum_difficulty(mut self, difficulty: u64) -> Self {
116 self.var_diff_config.minimum_difficulty = difficulty;
117 self
118 }
119
120 #[must_use]
121 pub fn with_maximum_difficulty(mut self, difficulty: u64) -> Self {
122 self.var_diff_config.maximum_difficulty = difficulty;
123 self
124 }
125
126 #[must_use]
127 pub fn with_retarget_time(mut self, time: u64) -> Self {
128 self.var_diff_config.retarget_time = time;
129 self
130 }
131
132 #[must_use]
133 pub fn with_target_time(mut self, time: u64) -> Self {
134 self.var_diff_config.target_time = time;
135 self
136 }
137
138 #[must_use]
139 pub fn with_variance_percent(mut self, percent: f64) -> Self {
140 self.var_diff_config.variance_percent = percent;
141 self
142 }
143
144 #[must_use]
145 pub fn with_initial_difficulty(mut self, difficulty: u64) -> Self {
146 self.var_diff_config.initial_difficulty = difficulty;
147 self
148 }
149
150 #[must_use]
151 pub fn with_ready_indicator(mut self, ready_indicator: ReadyIndicator) -> Self {
152 self.ready_indicator = ready_indicator;
153 self
154 }
155
156 #[must_use]
157 pub fn with_shutdown_message(mut self, msg: Buffer) -> Self {
158 self.shutdown_message = Some(msg);
159 self
160 }
161
162 #[must_use]
163 pub fn with_cancel_token(mut self, token: CancellationToken) -> Self {
164 self.cancel_token = Some(token);
165 self
166 }
167
168 #[must_use]
169 pub fn with_ban_manager(mut self, enabled: bool) -> Self {
170 self.ban_manager_enabled = enabled;
171 self
172 }
173
174 pub async fn build(self) -> Result<StratumServer<State, CState>> {
175 let ban_manager_config = BanManagerConfig {
176 enabled: self.ban_manager_enabled,
177 ..Default::default()
178 };
179
180 let config = Config {
181 connection: self.connection_config,
182 difficulty: self.var_diff_config,
183 bans: ban_manager_config,
184 };
185
186 let config_manager = ConfigManager::new(config);
187
188 let listener = TcpListener::bind(format!("{}:{}", self.host, self.port)).await?;
189
190 let listen_address = listener.local_addr()?;
192 let listener = TcpListenerStream::new(listener);
193 let session_list = SessionList::new(config_manager.clone());
194
195 let cancel_token = if let Some(cancel_token) = self.cancel_token {
196 cancel_token
197 } else {
198 CancellationToken::new()
199 };
200
201 let ban_manager = BanManager::new(config_manager.clone(), cancel_token.child_token());
202
203 #[cfg(feature = "api")]
204 let api = {
205 let state = crate::api::Context {
206 ban_manager: ban_manager.clone(),
207 ready_indicator: self.ready_indicator.create_new(),
208 };
209
210 let api_address = format!("{}:{}", self.api_host, self.api_port).parse()?;
211
212 crate::api::Api::build(api_address, state).await?
213 };
214
215 Ok(StratumServer {
216 id: self.server_id,
217 listener,
218 listen_address,
219 session_list,
220 config_manager,
221 state: self.state,
222 ban_manager,
223 router: Arc::new(Router::new()),
224 session_id_manager: IDManager::new(self.server_id),
225 cancel_token,
226 global_thread_list: JoinSet::new(),
227 ready_indicator: self.ready_indicator,
228 shutdown_message: self.shutdown_message,
229 #[cfg(feature = "api")]
230 api,
231 })
232 }
233}