From 3a0c19fbbf90b1fedcfce66daa110df88de80727 Mon Sep 17 00:00:00 2001
From: kamescg <kamescg@gmail.com>
Date: Mon, 28 Jun 2021 14:32:48 -0700
Subject: [PATCH] update to support SafeERC20

---
 contracts/BadgerYieldSource.sol | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/contracts/BadgerYieldSource.sol b/contracts/BadgerYieldSource.sol
index 495b701..780bd7b 100644
--- a/contracts/BadgerYieldSource.sol
+++ b/contracts/BadgerYieldSource.sol
@@ -3,6 +3,8 @@
 pragma solidity 0.6.12;
 
 import { IYieldSource } from "@pooltogether/yield-source-interface/contracts/IYieldSource.sol";
+import "@openzeppelin/contracts/token/ERC20/IERC20.sol";
+import "@openzeppelin/contracts/token/ERC20/SafeERC20.sol";
 import "@openzeppelin/contracts/math/SafeMath.sol";
 import "./IBadgerSett.sol";
 import "./IBadger.sol";
@@ -11,14 +13,15 @@ import "hardhat/console.sol";
 /// @title A pooltogether yield source for badger sett
 /// @author Steffel Fenix, 0xkarl
 contract BadgerYieldSource is IYieldSource {
+    using SafeERC20 for IERC20;
     using SafeMath for uint256;
     IBadgerSett private immutable badgerSett;
-    IBadger private immutable badger;
+    IERC20 private immutable badger;
     mapping(address => uint256) private balances;
 
-    constructor(address badgerSettAddr, address badgerAddr) public {
-        badgerSett = IBadgerSett(badgerSettAddr);
-        badger = IBadger(badgerAddr);
+    constructor(IBadgerSett badgerSettAddr, IERC20 badgerAddr) public {
+        badgerSett = badgerSettAddr;
+        badger = badgerAddr;
     }
 
     /// @notice Returns the ERC20 asset token used for deposits.
@@ -41,7 +44,7 @@ contract BadgerYieldSource is IYieldSource {
     /// @param amount The amount of `token()` to be supplied
     /// @param to The user whose balance will receive the tokens
     function supplyTokenTo(uint256 amount, address to) public override {
-        badger.transferFrom(msg.sender, address(this), amount);
+        badger.safeTransferFrom(msg.sender, address(this), amount);
         badger.approve(address(badgerSett), amount);
 
         uint256 beforeBalance = badgerSett.balanceOf(address(this));
@@ -76,7 +79,7 @@ contract BadgerYieldSource is IYieldSource {
         uint256 badgerBalanceDiff = badgerAfterBalance.sub(badgerBeforeBalance);
 
         balances[msg.sender] = balances[msg.sender].sub(requiredSharesBalance);
-        badger.transfer(msg.sender, badgerBalanceDiff);
+        badger.safeTransfer(msg.sender, badgerBalanceDiff);
         return (badgerBalanceDiff);
     }
 }