#961: Give each user its own schema for the lessons

This way we can reset a lesson using the database for each user and not for all users at once.
Also solves the issue that when someone solves the lesson it is solved for all users on the same WebGoat instance

webgoat-container/src/main/java/org/owasp/webgoat/DatabaseConfiguration.java

@@ -0,0 +1,56 @@
+package org.owasp.webgoat;
+
+import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.configuration.FluentConfiguration;
+import org.owasp.webgoat.service.RestartLessonService;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import javax.sql.DataSource;
+import java.util.Map;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+@Configuration
+public class DatabaseConfiguration {
+
+    private String driverClassName;
+
+    public DatabaseConfiguration(@Value("${spring.datasource.driver-class-name}") String driverClassName) {
+        this.driverClassName = driverClassName;
+    }
+
+    /**
+     * Define 2 Flyway instances, 1 for WebGoat itself which it uses for internal storage like users and 1 for lesson
+     * specific tables we use. This way we clean the data in the lesson database quite easily see {@link RestartLessonService#restartLesson()}
+     * for how we clean the lesson related tables.
+     */
+
+    @Bean(initMethod = "migrate")
+    public Flyway flyWayContainer(DataSource dataSource) {
+        return Flyway
+                .configure()
+                .configuration(Map.of("driver", driverClassName))
+                .dataSource(dataSource)
+                .schemas("container")
+                .locations("db/container")
+                .load();
+    }
+
+    @Bean
+    public Function<String, Flyway> flywayLessons(LessonDataSource lessonDataSource) {
+        return schema -> Flyway
+                .configure()
+                .configuration(Map.of("driver", driverClassName))
+                .schemas(schema)
+                .dataSource(lessonDataSource)
+                .load();
+    }
+
+    @Bean
+    public LessonDataSource lessonDataSource(DataSource dataSource) {
+        return new LessonDataSource(dataSource);
+    }
+}

webgoat-container/src/main/java/org/owasp/webgoat/DatabaseInitialization.java

@@ -1,50 +0,0 @@
-package org.owasp.webgoat;
-
-import org.flywaydb.core.Flyway;
-import org.owasp.webgoat.service.RestartLessonService;
-import org.springframework.beans.factory.annotation.Value;
-import org.springframework.context.annotation.Bean;
-import org.springframework.context.annotation.Configuration;
-import org.springframework.context.annotation.DependsOn;
-
-import javax.sql.DataSource;
-import java.util.Map;
-
-/**
- * Define 2 Flyway instances, 1 for WebGoat itself which it uses for internal storage like users and 1 for lesson
- * specific tables we use. This way we clean the data in the lesson database quite easily see {@link RestartLessonService#restartLesson()}
- * for how we clean the lesson related tables.
- */
-@Configuration
-public class DatabaseInitialization {
-
-    private final DataSource dataSource;
-    private String driverClassName;
-
-    public DatabaseInitialization(DataSource dataSource,
-                                  @Value("${spring.datasource.driver-class-name}") String driverClassName) {
-        this.dataSource = dataSource;
-        this.driverClassName = driverClassName;
-    }
-
-    @Bean(initMethod = "migrate")
-    public Flyway flyWayContainer() {
-        return Flyway
-                .configure()
-                .configuration(Map.of("driver", driverClassName))
-                .dataSource(dataSource)
-                .schemas("container")
-                .locations("db/container")
-                .load();
-    }
-
-    @Bean(initMethod = "migrate")
-    @DependsOn("flyWayContainer")
-    public Flyway flywayLessons() {
-        return Flyway
-                .configure()
-                .configuration(Map.of("driver", driverClassName))
-                .dataSource(dataSource)
-                .load();
-    }
-}

webgoat-container/src/main/java/org/owasp/webgoat/LessonDataSource.java

@@ -0,0 +1,70 @@
+package org.owasp.webgoat;
+
+import org.owasp.webgoat.lessons.LessonConnectionInvocationHandler;
+import org.springframework.jdbc.datasource.ConnectionProxy;
+
+import javax.sql.DataSource;
+import java.io.PrintWriter;
+import java.lang.reflect.Proxy;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.util.logging.Logger;
+
+public class LessonDataSource implements DataSource {
+
+    private final DataSource originalDataSource;
+
+    public LessonDataSource(DataSource dataSource) {
+        this.originalDataSource = dataSource;
+    }
+
+    @Override
+    public Connection getConnection() throws SQLException {
+        var targetConnection = originalDataSource.getConnection();
+        return (Connection) Proxy.newProxyInstance(
+                ConnectionProxy.class.getClassLoader(),
+                new Class[]{ConnectionProxy.class},
+                new LessonConnectionInvocationHandler(targetConnection));
+    }
+
+    @Override
+    public Connection getConnection(String username, String password) throws SQLException {
+        return originalDataSource.getConnection(username, password);
+    }
+
+    @Override
+    public PrintWriter getLogWriter() throws SQLException {
+        return originalDataSource.getLogWriter();
+    }
+
+    @Override
+    public void setLogWriter(PrintWriter out) throws SQLException {
+        originalDataSource.setLogWriter(out);
+    }
+
+    @Override
+    public void setLoginTimeout(int seconds) throws SQLException {
+        originalDataSource.setLoginTimeout(seconds);
+    }
+
+    @Override
+    public int getLoginTimeout() throws SQLException {
+        return originalDataSource.getLoginTimeout();
+    }
+
+    @Override
+    public Logger getParentLogger() throws SQLFeatureNotSupportedException {
+        return originalDataSource.getParentLogger();
+    }
+
+    @Override
+    public <T> T unwrap(Class<T> clazz) throws SQLException {
+        return originalDataSource.unwrap(clazz);
+    }
+
+    @Override
+    public boolean isWrapperFor(Class<?> clazz) throws SQLException {
+        return originalDataSource.isWrapperFor(clazz);
+    }
+}

webgoat-container/src/main/java/org/owasp/webgoat/lessons/LessonConnectionInvocationHandler.java

@@ -0,0 +1,36 @@
+package org.owasp.webgoat.lessons;
+
+import org.owasp.webgoat.users.WebGoatUser;
+import org.springframework.security.core.context.SecurityContextHolder;
+
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.sql.Connection;
+
+/**
+ * Handler which sets the correct schema for the currently bounded user. This way users are not seeing each other
+ * data and we can reset data for just one particular user.
+ */
+public class LessonConnectionInvocationHandler implements InvocationHandler {
+
+    private final Connection targetConnection;
+
+    public LessonConnectionInvocationHandler(Connection targetConnection) {
+        this.targetConnection = targetConnection;
+    }
+
+    @Override
+    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
+        var authentication = SecurityContextHolder.getContext().getAuthentication();
+        if (authentication != null && authentication.getPrincipal() instanceof WebGoatUser) {
+            var user = (WebGoatUser) authentication.getPrincipal();
+            targetConnection.createStatement().execute("SET SCHEMA \"" + user.getUsername() + "\"");
+        }
+        try {
+            return method.invoke(targetConnection, args);
+        } catch (InvocationTargetException e) {
+            throw e.getTargetException();
+        }
+    }
+}

webgoat-container/src/main/java/org/owasp/webgoat/service/RestartLessonService.java

@@ -36,6 +36,8 @@ import org.springframework.stereotype.Controller;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.ResponseStatus;
 
+import java.util.function.Function;
+
 @Controller
 @AllArgsConstructor
 @Slf4j
@@ -43,7 +45,7 @@ public class RestartLessonService {
 
     private final WebSession webSession;
     private final UserTrackerRepository userTrackerRepository;
-    private final Flyway flywayLessons;
+    private final Function<String, Flyway> flywayLessons;
 
     @RequestMapping(path = "/service/restartlesson.mvc", produces = "text/text")
     @ResponseStatus(value = HttpStatus.OK)
@@ -55,7 +57,8 @@ public class RestartLessonService {
         userTracker.reset(al);
         userTrackerRepository.save(userTracker);
 
-        flywayLessons.clean();
+        var flyway = flywayLessons.apply(webSession.getUserName());
-        flywayLessons.migrate();
+        flyway.clean();
+        flyway.migrate();
     }
 }

webgoat-container/src/main/java/org/owasp/webgoat/users/UserRepository.java

@@ -14,4 +14,6 @@ public interface UserRepository extends JpaRepository<WebGoatUser, String> {
 
     List<WebGoatUser> findAll();
 
+    boolean existsByUsername(String username);
+
 }

webgoat-container/src/main/java/org/owasp/webgoat/users/UserService.java

@@ -1,11 +1,16 @@
 package org.owasp.webgoat.users;
 
 import lombok.AllArgsConstructor;
+import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.configuration.FluentConfiguration;
+import org.owasp.webgoat.session.WebSession;
+import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.security.core.userdetails.UserDetailsService;
 import org.springframework.security.core.userdetails.UsernameNotFoundException;
 import org.springframework.stereotype.Service;
 
 import java.util.List;
+import java.util.function.Function;
 
 /**
  * @author nbaars
@@ -17,6 +22,8 @@ public class UserService implements UserDetailsService {
 
     private final UserRepository userRepository;
     private final UserTrackerRepository userTrackerRepository;
+    private final JdbcTemplate jdbcTemplate;
+    private final Function<String, Flyway> flywayLessons;
 
     @Override
     public WebGoatUser loadUserByUsername(String username) throws UsernameNotFoundException {
@@ -31,24 +38,18 @@ public class UserService implements UserDetailsService {
 
     public void addUser(String username, String password) {
         //get user if there exists one by the name
-        WebGoatUser webGoatUser = userRepository.findByUsername(username);
+        var userAlreadyExists = userRepository.existsByUsername(username);
-        //if user exists it will be updated, otherwise created
+        var webGoatUser = userRepository.save(new WebGoatUser(username, password));
-        userRepository.save(new WebGoatUser(username, password));
+
-        //if user previously existed it will not get another tracker
+        if (!userAlreadyExists) {
-        if (webGoatUser == null) {
+            userTrackerRepository.save(new UserTracker(username)); //if user previously existed it will not get another tracker
-            userTrackerRepository.save(new UserTracker(username));
+            createLessonsForUser(webGoatUser);
         }
     }
 
-    public void addUser(String username, String password, String role) {
+    private void createLessonsForUser(WebGoatUser webGoatUser) {
-        //get user if there exists one by the name
+        jdbcTemplate.execute("CREATE SCHEMA \"" + webGoatUser.getUsername() + "\" authorization dba");
-        WebGoatUser webGoatUser = userRepository.findByUsername(username);
+        flywayLessons.apply(webGoatUser.getUsername()).migrate();
-        //if user exists it will be updated, otherwise created
-        userRepository.save(new WebGoatUser(username, password, role));
-        //if user previously existed it will not get another tracker
-        if (webGoatUser == null) {
-            userTrackerRepository.save(new UserTracker(username));
-        }
     }
 
     public List<WebGoatUser> getAllUsers() {

webgoat-container/src/test/java/org/owasp/webgoat/TestApplication.java

@@ -5,6 +5,7 @@ import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Primary;
 import org.springframework.jdbc.datasource.DriverManagerDataSource;
 
 import javax.sql.DataSource;
@@ -14,14 +15,19 @@ import java.sql.SQLException;
 @SpringBootApplication
 public class TestApplication {
 
+    @Value("${spring.datasource.driver-class-name}")
+    private String driverClassName;
+
     /**
      * We define our own datasource, otherwise we end up with Hikari one which for some lessons will
      * throw an error (feature not supported)
      */
     @Bean
     @ConditionalOnProperty(prefix = "webgoat.start", name = "hsqldb", havingValue = "false")
+    @Primary
     public DataSource dataSource(@Value("${spring.datasource.url}") String url) throws SQLException {
         DriverManager.registerDriver(new JDBCDriver());
         return new DriverManagerDataSource(url);
     }
+
 }

webgoat-container/src/test/java/org/owasp/webgoat/plugins/LessonTest.java

@@ -1,5 +1,7 @@
 package org.owasp.webgoat.plugins;
 
+import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.configuration.FluentConfiguration;
 import org.junit.jupiter.api.BeforeEach;
 import org.owasp.webgoat.i18n.Language;
 import org.owasp.webgoat.i18n.PluginMessages;
@@ -12,7 +14,9 @@ import org.springframework.test.context.TestPropertySource;
 import org.springframework.test.web.servlet.MockMvc;
 import org.springframework.web.context.WebApplicationContext;
 
+import javax.annotation.PostConstruct;
 import java.util.Locale;
+import java.util.function.Function;
 
 import static org.mockito.Mockito.when;
 
@@ -31,6 +35,8 @@ public abstract class LessonTest {
     protected WebApplicationContext wac;
     @Autowired
     protected PluginMessages messages;
+    @Autowired
+    private Function<String, Flyway> flywayLessons;
     @MockBean
     protected WebSession webSession;
 
@@ -43,4 +49,11 @@ public abstract class LessonTest {
         when(language.getLocale()).thenReturn(Locale.getDefault());
     }
 
+    @PostConstruct
+    public void createFlywayLessonTables() {
+        flywayLessons.apply("PUBLIC").migrate();
+    }
+
+
+
 }

webgoat-container/src/test/java/org/owasp/webgoat/users/UserServiceTest.java

@@ -1,12 +1,18 @@
 package org.owasp.webgoat.users;
 
 import org.assertj.core.api.Assertions;
+import org.flywaydb.core.Flyway;
+import org.flywaydb.core.api.configuration.FluentConfiguration;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.Mock;
 import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.context.ApplicationEventPublisher;
+import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.security.core.userdetails.UsernameNotFoundException;
 
+import java.util.function.Function;
+
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.when;
 
@@ -17,11 +23,15 @@ class UserServiceTest {
     private UserRepository userRepository;
     @Mock
     private UserTrackerRepository userTrackerRepository;
+    @Mock
+    private JdbcTemplate jdbcTemplate;
+    @Mock
+    private Function<String, Flyway> flywayLessons;
 
     @Test
     void shouldThrowExceptionWhenUserIsNotFound() {
         when(userRepository.findByUsername(any())).thenReturn(null);
-        UserService userService = new UserService(userRepository, userTrackerRepository);
+        UserService userService = new UserService(userRepository, userTrackerRepository, jdbcTemplate, flywayLessons);
         Assertions.assertThatThrownBy(() -> userService.loadUserByUsername("unknown")).isInstanceOf(UsernameNotFoundException.class);
     }
 }

webgoat-lessons/challenge/src/main/java/org/owasp/webgoat/challenges/challenge5/Assignment5.java

@@ -23,6 +23,7 @@
 package org.owasp.webgoat.challenges.challenge5;
 
 import lombok.extern.slf4j.Slf4j;
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AttackResult;
 import org.owasp.webgoat.challenges.Flag;
@@ -40,9 +41,9 @@ import java.sql.ResultSet;
 @Slf4j
 public class Assignment5 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public Assignment5(DataSource dataSource) {
+    public Assignment5(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/jwt/src/main/java/org/owasp/webgoat/jwt/JWTFinalEndpoint.java

@@ -22,24 +22,15 @@
 
 package org.owasp.webgoat.jwt;
 
-import io.jsonwebtoken.Claims;
+import io.jsonwebtoken.*;
-import io.jsonwebtoken.JwsHeader;
-import io.jsonwebtoken.Jwt;
-import io.jsonwebtoken.JwtException;
-import io.jsonwebtoken.Jwts;
-import io.jsonwebtoken.SigningKeyResolverAdapter;
 import io.jsonwebtoken.impl.TextCodec;
 import org.apache.commons.lang3.StringUtils;
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
-import org.springframework.web.bind.annotation.PathVariable;
+import org.springframework.web.bind.annotation.*;
-import org.springframework.web.bind.annotation.PostMapping;
-import org.springframework.web.bind.annotation.RequestParam;
-import org.springframework.web.bind.annotation.ResponseBody;
-import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 
@@ -71,9 +62,9 @@ import java.sql.SQLException;
 @AssignmentHints({"jwt-final-hint1", "jwt-final-hint2", "jwt-final-hint3", "jwt-final-hint4", "jwt-final-hint5", "jwt-final-hint6"})
 public class JWTFinalEndpoint extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    private JWTFinalEndpoint(DataSource dataSource) {
+    private JWTFinalEndpoint(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
     

webgoat-lessons/missing-function-ac/src/main/java/org/owasp/webgoat/missing_ac/MissingFunctionACUsers.java

@@ -83,7 +83,7 @@ public class MissingFunctionACUsers {
     //@PreAuthorize()
     public WebGoatUser addUser(@RequestBody WebGoatUser newUser) {
         try {
-            userService.addUser(newUser.getUsername(),newUser.getPassword(),newUser.getRole());
+            userService.addUser(newUser.getUsername(),newUser.getPassword());
             return userService.loadUserByUsername(newUser.getUsername());
         } catch (Exception ex) {
             log.error("Error creating new User", ex);

webgoat-lessons/missing-function-ac/src/main/java/org/owasp/webgoat/missing_ac/Users.java

@@ -22,11 +22,11 @@
 
 package org.owasp.webgoat.missing_ac;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.session.UserSessionData;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.ResponseBody;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -36,9 +36,9 @@ import java.util.HashMap;
 public class Users {
 
     private UserSessionData userSessionData;
-    private DataSource dataSource;
+    private LessonDataSource dataSource;
 
-    public Users(UserSessionData userSessionData, DataSource dataSource) {
+    public Users(UserSessionData userSessionData, LessonDataSource dataSource) {
         this.userSessionData = userSessionData;
         this.dataSource = dataSource;
     }

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/advanced/SqlInjectionChallenge.java

@@ -23,6 +23,7 @@
 package org.owasp.webgoat.sql_injection.advanced;
 
 import lombok.extern.slf4j.Slf4j;
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -32,12 +33,7 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
+import java.sql.*;
-import java.sql.Connection;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.sql.Statement;
 
 /**
  * @author nbaars
@@ -48,9 +44,9 @@ import java.sql.Statement;
 @Slf4j
 public class SqlInjectionChallenge extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionChallenge(DataSource dataSource) {
+    public SqlInjectionChallenge(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/advanced/SqlInjectionChallengeLogin.java

@@ -22,6 +22,7 @@
 
 package org.owasp.webgoat.sql_injection.advanced;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -30,17 +31,13 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-
 @RestController
 @AssignmentHints(value = {"SqlInjectionChallengeHint1", "SqlInjectionChallengeHint2", "SqlInjectionChallengeHint3", "SqlInjectionChallengeHint4"})
 public class SqlInjectionChallengeLogin extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionChallengeLogin(DataSource dataSource) {
+    public SqlInjectionChallengeLogin(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/advanced/SqlInjectionLesson6a.java

@@ -22,6 +22,7 @@
 
 package org.owasp.webgoat.sql_injection.advanced;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,12 +32,7 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
+import java.sql.*;
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
-import java.sql.SQLException;
-import java.sql.Statement;
 
 
 @RestController
@@ -44,9 +40,9 @@ import java.sql.Statement;
         "SqlStringInjectionHint-advanced-6a-4", "SqlStringInjectionHint-advanced-6a-5"})
 public class SqlInjectionLesson6a extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson6a(DataSource dataSource) {
+    public SqlInjectionLesson6a(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/advanced/SqlInjectionLesson6b.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.advanced;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AttackResult;
 import org.springframework.web.bind.annotation.PostMapping;
@@ -30,7 +31,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.io.IOException;
 import java.sql.Connection;
 import java.sql.ResultSet;
@@ -41,9 +41,9 @@ import java.sql.Statement;
 @RestController
 public class SqlInjectionLesson6b extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson6b(DataSource dataSource) {
+    public SqlInjectionLesson6b(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson10.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -41,9 +41,9 @@ import java.sql.Statement;
 @AssignmentHints(value = {"SqlStringInjectionHint.10.1", "SqlStringInjectionHint.10.2", "SqlStringInjectionHint.10.3", "SqlStringInjectionHint.10.4", "SqlStringInjectionHint.10.5", "SqlStringInjectionHint.10.6"})
 public class SqlInjectionLesson10 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson10(DataSource dataSource) {
+    public SqlInjectionLesson10(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson2.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
@@ -44,9 +44,9 @@ import static java.sql.ResultSet.TYPE_SCROLL_INSENSITIVE;
 @AssignmentHints(value = {"SqlStringInjectionHint2-1", "SqlStringInjectionHint2-2", "SqlStringInjectionHint2-3", "SqlStringInjectionHint2-4"})
 public class SqlInjectionLesson2 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson2(DataSource dataSource) {
+    public SqlInjectionLesson2(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson3.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -45,9 +45,9 @@ import static java.sql.ResultSet.TYPE_SCROLL_INSENSITIVE;
 @AssignmentHints(value = {"SqlStringInjectionHint3-1", "SqlStringInjectionHint3-2"})
 public class SqlInjectionLesson3 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson3(DataSource dataSource) {
+    public SqlInjectionLesson3(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson4.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -45,9 +45,9 @@ import static java.sql.ResultSet.TYPE_SCROLL_INSENSITIVE;
 @AssignmentHints(value = {"SqlStringInjectionHint4-1", "SqlStringInjectionHint4-2", "SqlStringInjectionHint4-3"})
 public class SqlInjectionLesson4 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson4(DataSource dataSource) {
+    public SqlInjectionLesson4(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson5.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
 import javax.annotation.PostConstruct;
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -42,9 +42,9 @@ import java.sql.Statement;
 @AssignmentHints(value = {"SqlStringInjectionHint5-a"})
 public class SqlInjectionLesson5 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson5(DataSource dataSource) {
+    public SqlInjectionLesson5(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson5a.java

@@ -22,6 +22,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -30,12 +31,7 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
+import java.sql.*;
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
-import java.sql.SQLException;
-import java.sql.Statement;
 
 
 @RestController
@@ -46,9 +42,9 @@ public class SqlInjectionLesson5a extends AssignmentEndpoint {
             + "always evaluates to true (The string ending literal for '1 is closed by the query itself, so you should not inject it). "
             + "So the injected query basically looks like this: <span style=\"font-style: italic\">SELECT * FROM user_data WHERE first_name = 'John' and last_name = '' or TRUE</span>, "
             + "which will always evaluate to true, no matter what came before it.";
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson5a(DataSource dataSource) {
+    public SqlInjectionLesson5a(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson5b.java

@@ -22,6 +22,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,22 +32,17 @@ import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
 import javax.servlet.http.HttpServletRequest;
-import javax.sql.DataSource;
 import java.io.IOException;
-import java.sql.Connection;
+import java.sql.*;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
-import java.sql.SQLException;
 
 
 @RestController
 @AssignmentHints(value = {"SqlStringInjectionHint5b1", "SqlStringInjectionHint5b2", "SqlStringInjectionHint5b3", "SqlStringInjectionHint5b4"})
 public class SqlInjectionLesson5b extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson5b(DataSource dataSource) {
+    public SqlInjectionLesson5b(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson8.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,12 +32,7 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
+import java.sql.*;
-import java.sql.Connection;
-import java.sql.ResultSet;
-import java.sql.ResultSetMetaData;
-import java.sql.SQLException;
-import java.sql.Statement;
 import java.text.SimpleDateFormat;
 import java.util.Calendar;
 
@@ -47,9 +43,9 @@ import static java.sql.ResultSet.TYPE_SCROLL_SENSITIVE;
 @AssignmentHints(value = {"SqlStringInjectionHint.8.1", "SqlStringInjectionHint.8.2", "SqlStringInjectionHint.8.3", "SqlStringInjectionHint.8.4", "SqlStringInjectionHint.8.5"})
 public class SqlInjectionLesson8 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson8(DataSource dataSource) {
+    public SqlInjectionLesson8(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson9.java

@@ -23,6 +23,7 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
@@ -44,9 +44,9 @@ import static org.hsqldb.jdbc.JDBCResultSet.TYPE_SCROLL_SENSITIVE;
 @AssignmentHints(value = {"SqlStringInjectionHint.9.1", "SqlStringInjectionHint.9.2", "SqlStringInjectionHint.9.3", "SqlStringInjectionHint.9.4", "SqlStringInjectionHint.9.5"})
 public class SqlInjectionLesson9 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson9(DataSource dataSource) {
+    public SqlInjectionLesson9(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/mitigation/Servers.java

@@ -25,14 +25,10 @@ package org.owasp.webgoat.sql_injection.mitigation;
 import lombok.AllArgsConstructor;
 import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.owasp.webgoat.LessonDataSource;
 import org.springframework.http.MediaType;
-import org.springframework.web.bind.annotation.GetMapping;
+import org.springframework.web.bind.annotation.*;
-import org.springframework.web.bind.annotation.RequestMapping;
-import org.springframework.web.bind.annotation.RequestParam;
-import org.springframework.web.bind.annotation.ResponseBody;
-import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
@@ -48,7 +44,7 @@ import java.util.List;
 @Slf4j
 public class Servers {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
     @AllArgsConstructor
     @Getter
@@ -62,7 +58,7 @@ public class Servers {
         private String description;
     }
 
-    public Servers(DataSource dataSource) {
+    public Servers(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/main/java/org/owasp/webgoat/sql_injection/mitigation/SqlInjectionLesson13.java

@@ -23,6 +23,7 @@
 package org.owasp.webgoat.sql_injection.mitigation;
 
 import lombok.extern.slf4j.Slf4j;
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.assignments.AssignmentEndpoint;
 import org.owasp.webgoat.assignments.AssignmentHints;
 import org.owasp.webgoat.assignments.AttackResult;
@@ -31,7 +32,6 @@ import org.springframework.web.bind.annotation.RequestParam;
 import org.springframework.web.bind.annotation.ResponseBody;
 import org.springframework.web.bind.annotation.RestController;
 
-import javax.sql.DataSource;
 import java.sql.Connection;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
@@ -42,9 +42,9 @@ import java.sql.SQLException;
 @Slf4j
 public class SqlInjectionLesson13 extends AssignmentEndpoint {
 
-    private final DataSource dataSource;
+    private final LessonDataSource dataSource;
 
-    public SqlInjectionLesson13(DataSource dataSource) {
+    public SqlInjectionLesson13(LessonDataSource dataSource) {
         this.dataSource = dataSource;
     }
 

webgoat-lessons/sql-injection/src/test/java/org/owasp/webgoat/sql_injection/introduction/SqlInjectionLesson5Test.java

@@ -22,17 +22,16 @@
 
 package org.owasp.webgoat.sql_injection.introduction;
 
-import org.aspectj.lang.annotation.After;
 import org.hamcrest.CoreMatchers;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.owasp.webgoat.LessonDataSource;
 import org.owasp.webgoat.sql_injection.SqlLessonTest;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.test.context.junit.jupiter.SpringExtension;
 import org.springframework.test.web.servlet.request.MockMvcRequestBuilders;
 
-import javax.sql.DataSource;
 import java.sql.SQLException;
 
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath;
@@ -42,7 +41,7 @@ import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.
 public class SqlInjectionLesson5Test extends SqlLessonTest {
 
     @Autowired
-    private DataSource dataSource;
+    private LessonDataSource dataSource;
 
     @AfterEach
     public void removeGrant() throws SQLException {